"""Shared model setup pipeline for ModelZoo.

This module contains reusable setup mechanics that model-specific setup wrappers
can call from metadata/<MODEL>/setup.py.
"""

import json
import os
import subprocess
import sys
import tempfile
from pathlib import Path
from shutil import copy2, which
from typing import Any, Callable, Dict, List, Optional


HookMap = Dict[str, Callable[[Dict[str, Any]], None]]


def run_command(cmd: List[str], cwd: Optional[Path] = None, check: bool = True, env: Optional[Dict[str, str]] = None):
    """Run a command and stream output."""
    print(f"Running: {' '.join(cmd)}")
    result = subprocess.run(cmd, cwd=cwd, capture_output=False, env=env)
    if check and result.returncode != 0:
        raise RuntimeError(f"Command failed with exit code {result.returncode}: {' '.join(cmd)}")
    return result


def install_requirements_file(
    python_exe: Path,
    requirements_path: Path,
    filter_editable_git: bool = False,
    no_build_isolation: bool = False,
) -> int:
    """Install a single requirements file, optionally filtering editable git lines."""
    install_target = requirements_path
    temp_path: Optional[Path] = None
    try:
        if filter_editable_git:
            original_lines = requirements_path.read_text(encoding="utf-8").splitlines()
            filtered_lines = []
            for line in original_lines:
                stripped = line.strip()
                if stripped.startswith("-e git+"):
                    continue
                filtered_lines.append(line)
            if filtered_lines != original_lines:
                tmp = tempfile.NamedTemporaryFile(
                    mode="w",
                    encoding="utf-8",
                    suffix=".txt",
                    prefix="filtered_requirements_",
                    delete=False,
                )
                with tmp:
                    tmp.write("\n".join(filtered_lines))
                    tmp.write("\n")
                temp_path = Path(tmp.name)
                install_target = temp_path

        cmd = [str(python_exe), "-m", "pip", "install"]
        if no_build_isolation:
            cmd.append("--no-build-isolation")
        cmd.extend(["-r", str(install_target)])
        result = run_command(cmd, check=False)
        return result.returncode
    finally:
        if temp_path and temp_path.exists():
            temp_path.unlink()


def resolve_pyenv_base_command() -> List[str]:
    """Resolve a pyenv command that is executable from subprocess."""
    if sys.platform.startswith("win"):
        for name in ("pyenv.bat", "pyenv.cmd", "pyenv"):
            if which(name):
                return ["cmd", "/c", name]
    elif which("pyenv"):
        return ["pyenv"]
    return []


def run_pyenv(
    args: List[str],
    check: bool = True,
    capture_output: bool = False,
    text: bool = False,
    env: Optional[Dict[str, str]] = None,
):
    """Run pyenv in a cross-platform-safe way."""
    base = resolve_pyenv_base_command()
    if not base:
        raise RuntimeError(
            "pyenv was not found on PATH, but metadata requests a specific PythonVersion."
        )

    cmd = base + list(args)
    print(f"Running: {' '.join(cmd)}")
    try:
        return subprocess.run(
            cmd,
            capture_output=capture_output,
            text=text,
            check=check,
            env=env,
        )
    except FileNotFoundError as exc:
        raise RuntimeError(
            "pyenv appears configured in the shell but is not executable from subprocess. "
            "Ensure pyenv executable/shims are in PATH for non-interactive processes."
        ) from exc


def load_metadata(metadata_dir: Path) -> Dict[str, Any]:
    """Load model metadata from model.json or <model>.json."""
    candidates = [metadata_dir / "model.json", metadata_dir / f"{metadata_dir.name}.json"]
    for candidate in candidates:
        if candidate.exists():
            with candidate.open("r", encoding="utf-8") as handle:
                return json.load(handle)
    raise FileNotFoundError(
        f"Model metadata not found in {metadata_dir} "
        f"(expected model.json or {metadata_dir.name}.json)."
    )


def resolve_python_version(metadata: Dict[str, Any]) -> str:
    """Resolve requested Python version from metadata."""
    return (
        metadata.get("PythonVersion")
        or metadata.get("env", {}).get("pythonVersion")
        or ""
    )


def resolve_runtime_source(metadata: Dict[str, Any]) -> str:
    """Resolve architecture runtime source with backward-compatible defaults."""
    runtime_cfg = metadata.get("runtime") or {}
    runtime_source = str(runtime_cfg.get("source", "")).strip().lower()
    if runtime_source:
        return runtime_source

    # Backward compatibility for legacy metadata that only declared `repository`.
    if metadata.get("repository"):
        return "repository"
    return "standalone_script"


def resolve_patch_mode(metadata: Dict[str, Any], runtime_source: str) -> str:
    """Resolve patch mode defaults based on runtime source."""
    if "patchMode" in metadata:
        return str(metadata.get("patchMode") or "").strip().lower()
    if runtime_source == "repository":
        return "auto"
    return "none"


def pyenv_has_version(version: str) -> bool:
    """Check whether a Python version is already available in pyenv."""
    probe = run_pyenv(
        ["versions", "--bare"],
        capture_output=True,
        text=True,
        check=False,
    )
    if probe.returncode != 0:
        return False
    versions = {line.strip() for line in probe.stdout.splitlines() if line.strip()}
    return version in versions


def venv_python_path(venv_dir: Path) -> Path:
    """Return Python executable path for a virtualenv directory."""
    if sys.platform.startswith("win"):
        return venv_dir / "Scripts" / "python.exe"
    return venv_dir / "bin" / "python"


def ensure_venv(venv_dir: Path, python_version: str) -> Path:
    """Create venv if needed and return venv python path."""
    python_exe = venv_python_path(venv_dir)

    if python_exe.exists():
        print(f"Virtual environment already exists at {venv_dir}")
        return python_exe

    venv_dir.parent.mkdir(parents=True, exist_ok=True)

    if python_version:
        if not pyenv_has_version(python_version):
            run_pyenv(["install", python_version])

        pyenv_env = os.environ.copy()
        pyenv_env["PYENV_VERSION"] = python_version
        run_pyenv(
            ["exec", "python", "-m", "venv", str(venv_dir)],
            env=pyenv_env,
        )
    else:
        run_command([sys.executable, "-m", "venv", str(venv_dir)])

    if not python_exe.exists():
        raise RuntimeError(f"Virtual environment created, but python not found at {python_exe}")

    return python_exe


def install_dependencies(
    python_exe: Path,
    metadata_dir: Path,
    venv_dir: Path,
    requirements_files: Optional[List[str]] = None,
    fallback_file: str = "fallback_requirements.txt",
    local_repo_path: Optional[Path] = None,
    install_editable: bool = True,
) -> None:
    """Install requirements into the model-local venv."""
    if requirements_files is None:
        requirements_files = ["requirements_frozen.txt", "requirements_extra.txt"]

    # Keep setuptools in a range that still provides pkg_resources.
    run_command(
        [
            str(python_exe),
            "-m",
            "pip",
            "install",
            "--upgrade",
            "pip",
            "wheel",
            "setuptools<81",
        ]
    )

    for filename in requirements_files:
        path = metadata_dir / filename
        if not path.exists():
            print(f"Warning: {filename} not found in {metadata_dir}")
            continue

        result_code = install_requirements_file(
            python_exe=python_exe,
            requirements_path=path,
            filter_editable_git=local_repo_path is not None,
            no_build_isolation=False,
        )
        if result_code == 0:
            continue

        # Only frozen requirements get fallback logic.
        if filename != "requirements_frozen.txt":
            raise RuntimeError(f"Failed to install dependencies from {path}")

        print(
            "requirements_frozen.txt failed on first attempt. "
            "Retrying with --no-build-isolation for editable builds."
        )
        retry_code = install_requirements_file(
            python_exe=python_exe,
            requirements_path=path,
            filter_editable_git=local_repo_path is not None,
            no_build_isolation=True,
        )
        if retry_code == 0:
            continue

        fallback = metadata_dir / fallback_file
        if fallback.exists():
            print(
                "Frozen requirements still failed. "
                f"Falling back to {fallback_file}."
            )
            fallback_code = install_requirements_file(
                python_exe=python_exe,
                requirements_path=fallback,
                filter_editable_git=local_repo_path is not None,
                no_build_isolation=True,
            )
            if fallback_code != 0:
                raise RuntimeError(f"Failed to install fallback dependencies from {fallback}")
        else:
            raise RuntimeError(
                "requirements_frozen.txt failed and "
                f"{fallback_file} was not found."
            )

    if local_repo_path is not None and install_editable:
        editable_cmd = [str(python_exe), "-m", "pip", "install", "-e", str(local_repo_path)]
        editable_result = run_command(editable_cmd, check=False)
        if editable_result.returncode != 0:
            print(
                "Editable local repository install failed on first attempt. "
                "Retrying with --no-build-isolation."
            )
            editable_retry_cmd = [
                str(python_exe),
                "-m",
                "pip",
                "install",
                "--no-build-isolation",
                "-e",
                str(local_repo_path),
            ]
            editable_retry = run_command(editable_retry_cmd, check=False)
            if editable_retry.returncode != 0:
                raise RuntimeError(
                    "Failed to install local editable repository even with "
                    f"--no-build-isolation: {local_repo_path}"
                )
    elif local_repo_path is not None:
        print("Skipping editable install because repository.installEditable is false.")


def resolve_env_vars(raw_env: Dict[str, str], model_dir: Path, workspace_root: Path) -> Dict[str, str]:
    """Resolve ${MODEL_DIR} and relative values before download manager receives them."""
    resolved = {}
    for key, value in raw_env.items():
        if not isinstance(value, str):
            continue

        v = value.replace("${MODEL_DIR}", str(model_dir))
        if v.startswith("./") or v.startswith(".\\"):
            v = str(workspace_root / v[2:])

        resolved[key] = v
        os.environ[key] = v
        print(f"Set {key}={v}")
    return resolved


def copy_embeddings(metadata_dir: Path, model_dir: Path, embeddings_script: str = "embeddings.py") -> None:
    """Copy runtime scripts required for embeddings execution."""
    source = metadata_dir / embeddings_script
    if not source.exists():
        print(f"Warning: embeddings script not found at {source}")
        return

    target = model_dir / source.name
    copy2(source, target)
    print(f"Copied {source.name} -> {target}")

    shared_preprocess = Path(__file__).resolve().parent / "universal_templates" / "embeddings_preprocess.py"
    if not shared_preprocess.exists():
        print(f"Warning: shared preprocess script not found at {shared_preprocess}")
        return

    preprocess_target = model_dir / shared_preprocess.name
    copy2(shared_preprocess, preprocess_target)
    print(f"Copied {shared_preprocess.name} -> {preprocess_target}")


def apply_patches(
    metadata_dir: Path,
    target_root: Path,
    paths_filename: str = "paths.json",
) -> None:
    """Apply metadata patches into a target root directory."""
    patches_dir = metadata_dir / "patches"
    paths_json = patches_dir / paths_filename
    if not paths_json.exists():
        print("No patches/paths.json found; skipping patch step.")
        return

    if not target_root.exists():
        print(f"Warning: patch target path not found at {target_root}; skipping patches.")
        return

    with paths_json.open("r", encoding="utf-8") as handle:
        paths_map = json.load(handle)

    for patch_name, relative_target in paths_map.items():
        source = patches_dir / patch_name
        target = target_root / relative_target

        if not source.exists():
            print(f"Warning: patch file not found: {source}")
            continue

        target.parent.mkdir(parents=True, exist_ok=True)
        copy2(source, target)
        print(f"Patched {target}")


def ensure_repository_clone(repo_cfg: Dict[str, Any], model_dir: Path) -> Path:
    """Clone repository into model directory and checkout requested ref."""
    repo_url = repo_cfg.get("url")
    if not repo_url:
        raise RuntimeError("Repository configuration missing required key: url")

    requested_ref = repo_cfg.get("ref")
    repo_dir_name = repo_cfg.get("dir") or Path(repo_url).stem.replace(".git", "")
    repo_dir = model_dir / repo_dir_name

    if not repo_dir.exists():
        run_command(["git", "clone", repo_url, str(repo_dir)])
    elif not (repo_dir / ".git").exists():
        raise RuntimeError(f"Repository directory exists but is not a git repo: {repo_dir}")
    else:
        print(f"Repository already present: {repo_dir}")

    run_command(["git", "-C", str(repo_dir), "fetch", "--all", "--tags", "--quiet"], check=False)
    if requested_ref:
        run_command(["git", "-C", str(repo_dir), "checkout", "--quiet", requested_ref])
    return repo_dir


def resolve_repository_install_path(repo_dir: Path, repo_cfg: Dict[str, Any]) -> Path:
    """Resolve editable install path within a cloned repository."""
    raw_subdir = str(repo_cfg.get("installSubdir", ".")).strip()
    if not raw_subdir:
        raw_subdir = "."

    normalized = raw_subdir.replace("\\", "/")
    while normalized.startswith("./"):
        normalized = normalized[2:]
    while normalized.endswith("/"):
        normalized = normalized[:-1]

    install_dir = repo_dir if normalized in {"", "."} else repo_dir / normalized
    install_dir = install_dir.resolve()
    repo_root = repo_dir.resolve()

    try:
        install_dir.relative_to(repo_root)
    except ValueError:
        raise RuntimeError(
            "repository.installSubdir resolves outside repository root: "
            f"{raw_subdir}"
        )
    if not install_dir.exists():
        raise RuntimeError(
            "repository.installSubdir path does not exist: "
            f"{install_dir}"
        )
    return install_dir


def sentinel_path(model_dir: Path) -> Path:
    return model_dir / ".setup_complete"


def write_sentinel(model_dir: Path) -> None:
    path = sentinel_path(model_dir)
    path.write_text("setup_complete\n", encoding="utf-8")
    print(f"Setup complete. Sentinel written to {path}")


def is_setup_complete(model_dir: Path) -> bool:
    return sentinel_path(model_dir).exists()


def run_setup_pipeline(
    workspace_root: Path,
    metadata_dir: Path,
    model_name: Optional[str] = None,
    hooks: Optional[HookMap] = None,
    stage: str = "full",
) -> Dict[str, Any]:
    """Run shared setup pipeline for a model.

    Hooks are optional and keyed by step names:
    - pre_setup, post_setup
    - pre_venv, post_venv
    - pre_install, post_install
    - pre_download, post_download
    - pre_patches, post_patches
    """
    if hooks is None:
        hooks = {}
    if stage not in {"full", "download_only"}:
        raise ValueError(f"Unsupported setup stage: {stage}")

    try:
        from .downloads_manager import download_model_artifacts
    except ImportError:
        from downloads_manager import download_model_artifacts

    metadata = load_metadata(metadata_dir)
    resolved_model_name = model_name or metadata.get("modelName") or metadata_dir.name
    python_version = resolve_python_version(metadata)
    model_dir = workspace_root / "models" / resolved_model_name
    venv_dir = model_dir / ".venv"

    context = {
        "workspace_root": workspace_root,
        "metadata_dir": metadata_dir,
        "metadata": metadata,
        "model_name": resolved_model_name,
        "python_version": python_version,
        "model_dir": model_dir,
        "venv_dir": venv_dir,
        "python_exe": None,
        "repo_dir": None,
    }

    print(f"Setting up model: {resolved_model_name}")
    print(f"Workspace root: {workspace_root}")
    print(f"Model directory: {model_dir}")
    if python_version:
        print(f"Requested Python version: {python_version}")

    model_dir.mkdir(parents=True, exist_ok=True)

    if "pre_setup" in hooks:
        hooks["pre_setup"](context)

    raw_env = metadata.get("env", {}).get("variables", {})
    resolved_env = resolve_env_vars(raw_env, model_dir, workspace_root)
    context["resolved_env"] = resolved_env

    cache_file = metadata.get("cacheConfig", "model_cache.json")
    model_cache_path = metadata_dir / cache_file

    runtime_source = resolve_runtime_source(metadata)
    context["runtime_source"] = runtime_source
    repo_cfg = metadata.get("repository") if runtime_source == "repository" else None
    artifacts_in_repository = bool(repo_cfg.get("artifactsInRepository", False)) if isinstance(repo_cfg, dict) else False
    context["artifacts_in_repository"] = artifacts_in_repository

    if stage == "download_only":
        if "pre_download" in hooks:
            hooks["pre_download"](context)
        if artifacts_in_repository:
            if runtime_source == "repository":
                if not isinstance(repo_cfg, dict) or not repo_cfg.get("url"):
                    raise RuntimeError(
                        "repository.artifactsInRepository is true but repository.url is missing in metadata."
                    )
                context["repo_dir"] = ensure_repository_clone(repo_cfg, model_dir)
                context["repo_install_dir"] = resolve_repository_install_path(
                    context["repo_dir"],
                    repo_cfg,
                )
                print("Repository artifacts are bundled with source; ensured repository clone is present.")
            else:
                print("Skipping download phase because repository.artifactsInRepository is true.")
        elif model_cache_path.exists():
            download_model_artifacts(model_cache_path, workspace_root, env_vars=resolved_env)
        else:
            print(f"Warning: model cache config not found at {model_cache_path}")
        if "post_download" in hooks:
            hooks["post_download"](context)
        return context

    if "pre_venv" in hooks:
        hooks["pre_venv"](context)
    context["python_exe"] = ensure_venv(venv_dir, python_version)
    if "post_venv" in hooks:
        hooks["post_venv"](context)

    requirements_files = metadata.get("requirementsFiles")
    fallback_file = metadata.get("fallbackRequirementsFile", "fallback_requirements.txt")

    patch_mode = resolve_patch_mode(metadata, runtime_source)

    if runtime_source not in {"repository", "standalone_script"}:
        raise RuntimeError(
            "Unsupported metadata runtime.source: "
            f"{runtime_source}. Expected one of: repository, standalone_script"
        )

    install_editable = True
    if runtime_source == "repository":
        if not isinstance(repo_cfg, dict) or not repo_cfg.get("url"):
            raise RuntimeError(
                "runtime.source is 'repository' but repository.url is missing in metadata."
            )
        install_editable = bool(repo_cfg.get("installEditable", True))
        context["repo_dir"] = ensure_repository_clone(repo_cfg, model_dir)
        context["repo_install_dir"] = resolve_repository_install_path(
            context["repo_dir"],
            repo_cfg,
        )
    context["install_editable"] = install_editable

    if "pre_patches" in hooks:
        hooks["pre_patches"](context)
    if context["repo_dir"] is not None and patch_mode in ("auto", "pre_install"):
        apply_patches(metadata_dir, context["repo_dir"])

    if "pre_install" in hooks:
        hooks["pre_install"](context)
    install_dependencies(
        context["python_exe"],
        metadata_dir,
        venv_dir,
        requirements_files=requirements_files,
        fallback_file=fallback_file,
        local_repo_path=context.get("repo_install_dir", context["repo_dir"]),
        install_editable=bool(context.get("install_editable", True)),
    )
    if "post_install" in hooks:
        hooks["post_install"](context)

    embeddings_script = metadata.get("embeddingsScript", "embeddings.py")
    copy_embeddings(metadata_dir, model_dir, embeddings_script=embeddings_script)

    if "pre_download" in hooks:
        hooks["pre_download"](context)
    if artifacts_in_repository:
        print("Skipping download phase because repository.artifactsInRepository is true.")
    elif model_cache_path.exists():
        download_model_artifacts(model_cache_path, workspace_root, env_vars=resolved_env)
    else:
        print(f"Warning: model cache config not found at {model_cache_path}")
    if "post_download" in hooks:
        hooks["post_download"](context)

    if context["repo_dir"] is None and patch_mode in ("auto", "post_install"):
        patch_target_subdir = metadata.get("patchTargetDir", resolved_model_name.lower())
        target_root = venv_dir / "src" / patch_target_subdir
        apply_patches(metadata_dir, target_root)
    if "post_patches" in hooks:
        hooks["post_patches"](context)

    write_sentinel(model_dir)
    if "post_setup" in hooks:
        hooks["post_setup"](context)

    return context
