"""
llmpm managed virtualenv.

All llmpm commands run inside ~/.llmpm/venv to guarantee a clean, isolated
package set regardless of system-level Python installs (Homebrew, conda, etc.).

Usage
-----
    from llmpm._venv import ensure
    ensure()   # no-op if already in the venv; otherwise creates + re-execs

Set LLPM_NO_VENV=1 to bypass (useful in CI, Docker, or containers where
isolation is already provided by the environment).
"""
from __future__ import annotations

import hashlib
import json
import os
import pathlib
import shutil
import subprocess
import sys
import threading
import time

# ── paths ─────────────────────────────────────────────────────────────────────

VENV_DIR = pathlib.Path.home() / ".llmpm" / "venv"
_IS_WIN = sys.platform == "win32"
_BIN = "Scripts" if _IS_WIN else "bin"
_PY_EXE = "python.exe" if _IS_WIN else "python"
_PIP_EXE = "pip.exe" if _IS_WIN else "pip"

VENV_PYTHON: pathlib.Path = VENV_DIR / _BIN / _PY_EXE
VENV_PIP: pathlib.Path = VENV_DIR / _BIN / _PIP_EXE

# Tracks what was installed so we can detect when deps have changed.
_MARKER: pathlib.Path = VENV_DIR / ".llmpm_version"

# ── ANSI helpers (stdlib-only, works before rich is available) ─────────────────

_RESET = "\x1b[0m"
_BOLD  = "\x1b[1m"
_DIM   = "\x1b[2m"
_GREEN = "\x1b[32m"
_RED   = "\x1b[31m"

# ASCII fallbacks for Windows CMD, which renders UTF-8 bytes as cp1252 garbage.
_CHECK = "+" if _IS_WIN else "✓"
_CROSS = "x" if _IS_WIN else "✗"
_SEP   = ("-" if _IS_WIN else "─") * 54


def _enable_windows_vt() -> None:
    """Enable ANSI escape processing on Windows 10+ (ENABLE_VIRTUAL_TERMINAL_PROCESSING)."""
    if not _IS_WIN:
        return
    try:
        import ctypes
        import ctypes.wintypes
        k32 = ctypes.windll.kernel32
        handle = k32.GetStdHandle(-11)  # STD_OUTPUT_HANDLE
        mode = ctypes.wintypes.DWORD()
        k32.GetConsoleMode(handle, ctypes.byref(mode))
        k32.SetConsoleMode(handle, mode.value | 0x0004)  # ENABLE_VIRTUAL_TERMINAL_PROCESSING
    except Exception:
        pass


_enable_windows_vt()


def _print(msg: str = "") -> None:
    os.write(1, (msg + "\n").encode())


def _ok(msg: str) -> None:
    _print(f"  {_GREEN}{_CHECK}{_RESET}  {msg}")


def _fail(msg: str) -> None:
    _print(f"  {_RED}{_CROSS}{_RESET}  {msg}")


def _step(n: int, total: int, msg: str) -> None:
    _print(f"\n  {_BOLD}Step {n}/{total}{_RESET}  {msg}\n")


# ── pip progress display ───────────────────────────────────────────────────────

_SPIN  = (["-", "\\", "|", "/"] if _IS_WIN
          else ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"])
_FILL  = "#" if _IS_WIN else "█"
_EMPTY = "." if _IS_WIN else "░"
_BAR_W = 34


def _bar(n: int, total: int) -> str:
    if total == 0:
        return _EMPTY * _BAR_W
    filled = min(_BAR_W, int(_BAR_W * n / total))
    return f"{_GREEN}{_FILL * filled}{_RESET}{_EMPTY * (_BAR_W - filled)}"


def _over(msg: str) -> None:
    """Overwrite the current terminal line in-place."""
    os.write(1, f"\r  {msg:<77}".encode())


def _run_pip(cmd: list[str], **kwargs) -> subprocess.CompletedProcess:
    """
    Run a pip command with a live animated progress display.

    A background thread redraws the current line every 100 ms so the spinner
    keeps moving even when pip produces no output (e.g. the silent install
    phase after all packages have been downloaded).

    Phases: Resolving → Downloading → Building → Installing → Done
    """
    # Force unbuffered output from pip.
    # When stdout is a pipe (not a TTY) Python switches to full block-buffering,
    # so all lines arrive at once at the end instead of streaming line-by-line.
    # -u disables that at the interpreter level; PYTHONUNBUFFERED covers any
    # subprocess that pip itself might spawn (e.g. build backends).
    unbuf_cmd = [cmd[0], "-u", *cmd[1:]]
    unbuf_env = {**os.environ, "PYTHONUNBUFFERED": "1"}

    proc = subprocess.Popen(
        unbuf_cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
        env=unbuf_env,
        **kwargs,
    )
    assert proc.stdout is not None

    # Shared state written by the reader (main thread), read by animator.
    # CPython's GIL makes individual dict-value writes atomic enough here.
    state: dict = {
        "phase": "resolve",
        "collected": 0,
        "install_total": 0,
        "install_pkgs": [],   # package names parsed from "Installing collected packages:"
        "install_idx": 0,     # index of package currently shown in spinner
        "current": "",
    }
    stop = threading.Event()

    def _render(tick: int) -> None:
        s = _SPIN[tick % len(_SPIN)]
        p = state["phase"]
        if p == "resolve":
            _over(
                f"{_DIM}Resolving   {_RESET}{s}  "
                f"{state['collected']} packages  {_DIM}{state['current'][:30]}{_RESET}"
            )
        elif p == "download":
            _over(f"{_DIM}Downloading {_RESET}{s}  {_DIM}{state['current'][:40]}{_RESET}")
        elif p == "build":
            _over(f"{_DIM}Building    {_RESET}{s}  {_DIM}compiling {state['current'][:30]}{_RESET}")
        elif p == "install":
            done  = state["install_idx"]
            total = state["install_total"]
            pkgs  = state["install_pkgs"]
            name  = pkgs[done] if done < len(pkgs) else ""
            _over(
                f"{_DIM}Installing  {_RESET}{s}  "
                f"[{_bar(done, total)}]  {_DIM}{done}/{total}{_RESET}  "
                f"\x1b[36m{name[:32]}{_RESET}"
            )

    def _animate() -> None:
        tick = 0
        while not stop.wait(0.1):   # redraws every 100 ms
            # Advance the displayed package name every ~500 ms during install.
            # pip doesn't emit per-package lines, so we cycle through the list
            # as a visual indicator of progress.
            if state["phase"] == "install" and tick % 5 == 0:
                n = state["install_total"]
                if n > 0:
                    state["install_idx"] = min(state["install_idx"] + 1, n - 1)
            _render(tick)
            tick += 1

    def _start() -> threading.Thread:
        stop.clear()
        t = threading.Thread(target=_animate, daemon=True)
        t.start()
        return t

    def _stop(t: threading.Thread) -> None:
        stop.set()
        t.join()

    anim = _start()
    errors: list[str] = []
    final_phase = "resolve"

    for raw in proc.stdout:
        line = raw.rstrip()
        if not line:
            continue

        if line.startswith("Collecting "):
            state["collected"] += 1
            state["current"] = line.split()[1]
            state["phase"] = "resolve"

        elif line.startswith("  Downloading ") or line.startswith("Downloading "):
            state["current"] = line.strip().split()[1].split("/")[-1]
            state["phase"] = "download"

        elif "Building wheel" in line or "building wheel" in line:
            if "for" in line:
                state["current"] = line.split("for")[-1].strip().split()[0]
            state["phase"] = "build"

        elif line.startswith("Installing collected packages:"):
            pkgs_str = line[len("Installing collected packages:"):].strip().rstrip(".")
            pkg_list = [p.strip() for p in pkgs_str.split(",") if p.strip()]
            state["install_pkgs"] = pkg_list
            state["install_total"] = len(pkg_list)
            state["install_idx"] = 0
            _stop(anim)
            os.write(1, b"\n")          # commit the resolving/download line
            state["phase"] = "install"
            anim = _start()             # restart animator for install phase

        elif line.startswith("Successfully installed"):
            final_phase = "done"
            installed = len(line.split()) - 2
            _stop(anim)
            os.write(1, b"\r" + b" " * 80 + b"\r")
            _ok(f"{installed} packages installed successfully")

        elif line.upper().startswith("ERROR") or "ERROR:" in line:
            errors.append(line)
            _stop(anim)
            os.write(1, b"\r" + b" " * 80 + b"\r")
            _fail(line)
            os.write(1, b"\n")
            anim = _start()

        # suppress routine pip noise (warnings, "Already satisfied", etc.)

    proc.wait()
    _stop(anim)

    if final_phase != "done":
        os.write(1, b"\n")
        for err in errors:
            _fail(err)

    return subprocess.CompletedProcess(cmd, proc.returncode)


# ── helpers ───────────────────────────────────────────────────────────────────

def in_venv() -> bool:
    """Return True if the current Python is running inside VENV_DIR.

    Uses sys.prefix rather than sys.executable because on macOS with Homebrew
    the venv Python binary is a symlink chain that resolves outside the venv
    directory, causing is_relative_to(VENV_DIR) to return False even when
    running inside the venv.  sys.prefix is set explicitly by the venv module
    to point at the venv root and does not involve symlink resolution.
    """
    try:
        venv_path = VENV_DIR.resolve()
    except OSError:
        venv_path = VENV_DIR.absolute()
    try:
        prefix_path = pathlib.Path(sys.prefix).resolve()
    except OSError:
        prefix_path = pathlib.Path(sys.prefix).absolute()
    try:
        return prefix_path == venv_path or prefix_path.is_relative_to(venv_path)
    except AttributeError:  # Python < 3.9 has no is_relative_to
        venv_str = str(venv_path)
        prefix_str = str(prefix_path)
        return prefix_str == venv_str or prefix_str.startswith(venv_str + os.sep)


def _find_editable_root() -> pathlib.Path | None:
    """
    Return the project root if this file lives inside a dev checkout
    (i.e. a directory that contains pyproject.toml).
    """
    candidate = pathlib.Path(__file__).resolve().parent.parent
    if (candidate / "pyproject.toml").exists():
        return candidate
    return None


# ── version / dependency tracking ─────────────────────────────────────────────

def _current_version() -> str:
    """Return the currently running llmpm version string."""
    try:
        from importlib.metadata import version as _iv
        return _iv("llmpm")
    except Exception:
        pass
    try:
        init = pathlib.Path(__file__).resolve().parent / "__init__.py"
        for line in init.read_text().splitlines():
            if line.startswith("__version__"):
                return line.split('"')[1]
    except Exception:
        pass
    return "unknown"


def _pyproject_hash() -> str:
    """Return an MD5 hash of pyproject.toml content if in a dev checkout."""
    root = _find_editable_root()
    if root is not None:
        p = root / "pyproject.toml"
        if p.exists():
            try:
                return hashlib.md5(p.read_bytes()).hexdigest()
            except OSError:
                pass
    return ""


def _write_marker() -> None:
    """Stamp the venv with the current version + pyproject.toml hash."""
    _MARKER.write_text(json.dumps({
        "version": _current_version(),
        "pyproject_hash": _pyproject_hash(),
    }))


def _needs_rebuild() -> bool:
    """Return True if the venv is out of date with the installed llmpm."""
    try:
        marker = json.loads(_MARKER.read_text())
    except Exception:
        return True
    # Migrate from old mtime-based marker format → force one rebuild
    if "pyproject_mtime" in marker:
        return True
    if marker.get("version") != _current_version():
        return True
    # Dev mode: rebuild when pyproject.toml content changes.
    # Compare directly (no empty-string guard) so that a transition from a
    # previously-undetected root ("") to a real hash also triggers a rebuild.
    if _pyproject_hash() != marker.get("pyproject_hash", ""):
        return True
    return False


# ── venv creation ─────────────────────────────────────────────────────────────

def _create() -> None:
    """Create VENV_DIR and install llmpm (+ all backends) into it."""
    import venv as _venv_mod  # stdlib

    _print()
    _print(f"  {_BOLD}llmpm{_RESET}  Setting up isolated environment")
    _print(f"  {_SEP}")
    _print(f"  {_DIM}Location: {VENV_DIR}{_RESET}")
    _print()

    # ── Step 1: create the virtualenv ─────────────────────────────────────────
    _step(1, 3, "Creating virtual environment …")
    _venv_mod.create(str(VENV_DIR), with_pip=True, clear=True)
    _ok(f"{VENV_DIR}")

    # ── Step 2: upgrade pip ───────────────────────────────────────────────────
    _step(2, 3, "Upgrading pip …")
    r = _run_pip([str(VENV_PYTHON), "-m", "pip", "install", "--upgrade", "pip"])
    if r.returncode == 0:
        _ok("pip up to date")
    else:
        _fail("pip upgrade failed (continuing anyway)")

    editable_root = _find_editable_root()

    # ── Step 3: install llmpm core (must succeed) ─────────────────────────────
    # Install only the lightweight CLI deps first so the venv is always
    # usable even if the heavy ML backends fail in step 4.
    if editable_root is not None:
        _step(3, 4, f"Installing llmpm core (editable) from {editable_root} …")
        r = _run_pip([str(VENV_PYTHON), "-m", "pip", "install", "-e", str(editable_root)])
    else:
        _step(3, 4, "Installing llmpm core from PyPI …")
        r = _run_pip([str(VENV_PYTHON), "-m", "pip", "install", "llmpm"])

    if r.returncode != 0:
        _print()
        _print(f"  {_SEP}")
        _fail("Core installation failed — see pip output above.")
        _print(f"  {_DIM}Try running:  pip install llmpm{_RESET}")
        _print()
        raise RuntimeError("llmpm core installation into venv failed")
    _ok("llmpm core installed")

    # ── Step 4: install ML backends (best-effort) ─────────────────────────────
    # Heavy backends (torch, llama-cpp-python, etc.) may fail on some platforms;
    # that is non-fatal — the CLI still works and backends can be added later.
    if editable_root is not None:
        _step(4, 4, f"Installing ML backends (editable [all]) …")
        r = _run_pip([
            str(VENV_PYTHON), "-m", "pip", "install",
            "-e", f"{editable_root}[all]",
        ])
    else:
        _step(4, 4, "Installing ML backends (torch, transformers, llama-cpp-python…) …")
        # Read [all] extra deps directly from the installed llmpm metadata and
        # install them as plain packages.  Avoids pip re-resolving llmpm[all]
        # across every historical PyPI version (which fails on Python 3.13+
        # because old versions pinned torch<2.5 / numpy<2.0).
        try:
            from importlib.metadata import requires as _requires
            _all_deps = [
                req.split(";")[0].strip()
                for req in (_requires("llmpm") or [])
                if 'extra == "all"' in req
            ]
        except Exception:
            _all_deps = []

        if not _all_deps:
            # Fallback: pin to exact installed version to avoid cross-version conflict.
            _ver = _current_version()
            _spec = f"llmpm[all]=={_ver}" if _ver != "unknown" else "llmpm[all]"
            r = _run_pip([str(VENV_PYTHON), "-m", "pip", "install", _spec])
        else:
            # Install in dependency order to avoid resolution conflicts:
            #   1. Standalone deps (no torch required)
            #   2. torch + llama-cpp-python (platform-sensitive; try unconstrained on failure)
            #   3. torch-dependent packages (only if torch succeeded)
            _STANDALONE  = ("numpy", "Pillow", "soundfile")
            _TORCH_FIRST = ("torch", "llama-cpp-python")
            _TORCH_DEP   = ("transformers", "accelerate", "diffusers")

            _standalone = [d for d in _all_deps if any(p in d for p in _STANDALONE)]
            _torch_pkgs = [d for d in _all_deps if any(p in d for p in _TORCH_FIRST)]
            _torch_deps = [d for d in _all_deps if any(p in d for p in _TORCH_DEP)]

            # Step 1 — standalone (always succeeds on supported Python)
            r = (_run_pip([str(VENV_PYTHON), "-m", "pip", "install", *_standalone])
                 if _standalone else type("R", (), {"returncode": 0})())

            # Step 2 — torch + llama-cpp-python
            _torch_ok = True
            if _torch_pkgs:
                r2 = _run_pip([str(VENV_PYTHON), "-m", "pip", "install", *_torch_pkgs])
                if r2.returncode != 0:
                    # Retry without version constraints (lets pip pick the newest wheel)
                    _unconstrained = [p.split(">")[0].split("<")[0].split("=")[0]
                                      for p in _torch_pkgs]
                    r2 = _run_pip(
                        [str(VENV_PYTHON), "-m", "pip", "install", *_unconstrained]
                    )
                if r2.returncode != 0:
                    _torch_ok = False
                    _print()
                    _fail("torch / llama-cpp-python could not be installed automatically.")
                    _print(
                        f"  {_DIM}This is common on Python 3.13 or unusual platforms.{_RESET}"
                    )
                    _print(
                        f"  {_DIM}Install manually:  pip install torch"
                        f"  (or visit https://pytorch.org/get-started){_RESET}"
                    )

            # Step 3 — torch-dependent packages (skip entirely if torch failed)
            if _torch_ok and _torch_deps:
                _run_pip([str(VENV_PYTHON), "-m", "pip", "install", *_torch_deps])

    _print()
    _print(f"  {_SEP}")
    if r.returncode != 0:
        _fail("ML backends install failed — some inference features may be unavailable.")
        _print(f"  {_DIM}Retry with:  pip install \"llmpm[all]\"{_RESET}")
    else:
        _ok(f"Environment ready  {_DIM}({VENV_DIR}){_RESET}")
    _write_marker()
    _print()


def _rebuild() -> None:
    """Wipe and recreate the venv because llmpm or its deps have changed."""
    _print()
    _print(f"  {_BOLD}llmpm{_RESET}  Dependency change detected — rebuilding environment")
    _print(f"  {_SEP}")
    _print(f"  {_DIM}Removing: {VENV_DIR}{_RESET}")
    _print()
    shutil.rmtree(str(VENV_DIR), ignore_errors=True)
    _create()


# ── public API ────────────────────────────────────────────────────────────────

def ensure() -> None:
    """
    Guarantee that llmpm runs inside its managed virtualenv.

    * Already in the venv      → returns immediately (no-op).
    * LLPM_NO_VENV=1           → returns immediately (bypass for CI/Docker).
    * Venv missing             → creates it, installs llmpm[all], then re-execs.
    * Venv present + outdated  → wipes it, recreates it, then re-execs.
    * Venv present + current   → re-execs with the venv Python.

    The re-exec uses os.execv(), which *replaces* the current process
    (no subprocess overhead, no duplicate output).
    """
    if os.environ.get("LLPM_NO_VENV"):
        return
    if in_venv():
        return
    if not VENV_PYTHON.exists():
        try:
            _create()
        except Exception as exc:
            _fail(f"Environment setup failed ({exc}) — retrying from scratch …")
            shutil.rmtree(str(VENV_DIR), ignore_errors=True)
            _create()
    elif _needs_rebuild():
        try:
            _rebuild()
        except Exception as exc:
            _fail(f"Rebuild failed ({exc}) — retrying from scratch …")
            shutil.rmtree(str(VENV_DIR), ignore_errors=True)
            _create()
    # Hand off to the venv Python running the same llmpm command
    os.execv(str(VENV_PYTHON), [str(VENV_PYTHON), "-m", "llmpm", *sys.argv[1:]])


def pip_install(*packages: str, quiet: bool = False) -> int:
    """
    Install *packages* into the llmpm venv.

    Returns pip's exit code (0 = success).
    Useful for commands that want to install optional backends into the venv.
    """
    cmd = [str(VENV_PYTHON), "-m", "pip", "install"]
    if quiet:
        cmd.append("--quiet")
    cmd.extend(packages)
    return subprocess.run(cmd, check=False).returncode
