#!/usr/bin/env python3
import argparse
import json
import pathlib
import sys

from script_common import PROXY_VARS, clean_proxy_env, dotnet_env, ensure_dirs, pip_env, resolve_repo_root


def huggingface_env(repo_root: pathlib.Path) -> dict[str, str]:
    env = {}
    hf_home = repo_root / ".cache" / "huggingface"
    hf_hub_cache = hf_home / "hub"
    ensure_dirs([hf_hub_cache])
    env["HF_HOME"] = str(hf_home)
    env["HUGGINGFACE_HUB_CACHE"] = str(hf_hub_cache)
    env["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
    return env


def resolved_env(repo_root: pathlib.Path) -> dict[str, str]:
    env = {}
    dotnet = dotnet_env(repo_root)
    pip = pip_env(repo_root)
    hf = huggingface_env(repo_root)

    dotnet_keys = [
        "DOTNET_CLI_HOME",
        "NUGET_PACKAGES",
        "NUGET_HTTP_CACHE_PATH",
        "DOTNET_SKIP_FIRST_TIME_EXPERIENCE",
        "DOTNET_ADD_GLOBAL_TOOLS_TO_PATH",
        "DOTNET_GENERATE_ASPNET_CERTIFICATE",
        "DOTNET_CLI_TELEMETRY_OPTOUT",
        "NUGET_CERT_REVOCATION_MODE",
    ]
    pip_keys = [
        "PIP_CACHE_DIR",
        "PIP_DISABLE_PIP_VERSION_CHECK",
        "PIP_DEFAULT_TIMEOUT",
        "PIP_RETRIES",
        "TEMP",
        "TMP",
    ]
    for key in dotnet_keys:
        env[key] = dotnet[key]
    for key in pip_keys:
        env[key] = pip[key]
    env.update(hf)
    clean_proxy_env(env)
    return env


def export_lines(shell: str, env_map: dict[str, str]) -> list[str]:
    def sh_quote(value: str) -> str:
        return "'" + value.replace("'", "'\"'\"'") + "'"

    if shell == "pwsh":
        lines = [f"Remove-Item Env:{k} -ErrorAction SilentlyContinue" for k in PROXY_VARS]
        lines.extend(f"$env:{k} = \"{v.replace('\"', '`\"')}\"" for k, v in env_map.items())
        return lines
    if shell in ("bash", "zsh"):
        lines = [f"unset {k}" for k in PROXY_VARS]
        lines.extend(f"export {k}={sh_quote(v)}" for k, v in env_map.items())
        return lines
    if shell == "cmd":
        lines = [f"set {k}=" for k in PROXY_VARS]
        lines.extend(f"set {k}={v}" for k, v in env_map.items())
        return lines
    raise ValueError(shell)


def cmd_export(args):
    try:
        repo_root = resolve_repo_root(args.project_root)
    except Exception as ex:
        print(f"Failed to resolve project root: {ex}", file=sys.stderr)
        return 2

    env_map = resolved_env(repo_root)
    payload = {
        "projectRoot": str(repo_root),
        "env": env_map,
        "createdDirs": [
            str(repo_root / ".dotnet_home"),
            str(repo_root / ".nuget" / "packages"),
            str(repo_root / ".nuget" / "http-cache"),
            str(repo_root / ".pip" / "cache"),
            str(repo_root / ".tmp" / "pip-temp"),
            str(repo_root / ".cache" / "huggingface" / "hub"),
        ],
        "warnings": [],
    }

    try:
        lines = export_lines(args.shell, env_map)
    except ValueError:
        print(f"Unsupported shell target: {args.shell}", file=sys.stderr)
        return 3

    if args.json:
        print(json.dumps(payload))
    else:
        for line in lines:
            print(line)
    return 0


def cmd_doctor(args):
    try:
        repo_root = resolve_repo_root(args.project_root)
    except Exception as ex:
        print(f"Failed to resolve project root: {ex}", file=sys.stderr)
        return 2

    env_map = resolved_env(repo_root)
    checks = {
        "repo_root": str(repo_root),
        "dotnet_home_exists": (repo_root / ".dotnet_home").exists(),
        "nuget_cache_exists": (repo_root / ".nuget" / "packages").exists(),
        "pip_cache_exists": (repo_root / ".pip" / "cache").exists(),
        "hf_cache_exists": (repo_root / ".cache" / "huggingface" / "hub").exists(),
        "env_count": len(env_map),
    }
    print(json.dumps(checks))
    return 0


def main():
    parser = argparse.ArgumentParser(description="SDT cross-platform shell bootstrap helper")
    sub = parser.add_subparsers(dest="command", required=True)

    p_export = sub.add_parser("export", help="Print env exports for a shell")
    p_export.add_argument("--shell", required=True)
    p_export.add_argument("--project-root")
    p_export.add_argument("--json", action="store_true")

    p_doctor = sub.add_parser("doctor", help="Validate env bootstrap paths")
    p_doctor.add_argument("--project-root")

    args = parser.parse_args()
    if args.command == "export":
        return cmd_export(args)
    return cmd_doctor(args)


if __name__ == "__main__":
    sys.exit(main())