setup/fix_install.py

"""Diagnose and repair the import failures left by setup/install.py.

Failures from the previous run:
  1. seqeval     -> build-time error (setuptools_scm 10.x bug: vcs_versioning)
  2. gradio      -> imports installed package but fails on import (likely starlette 1.0)
  3. fastapi     -> same root cause as gradio (starlette 1.0 broke compat)
  4. wandb       -> optional; protobuf 7 may be the cause

Strategy:
  Step 1: print the FULL import traceback for each failing package so we know
          exactly what is wrong (not just the exception class name).
  Step 2: apply targeted fixes:
            - seqeval: install setuptools_scm<10 first, then seqeval
            - fastapi/gradio: pin starlette<1.0 + known-good versions
            - wandb (optional): pin protobuf<6 + reinstall
  Step 3: re-verify all four imports.
  Step 4: rewrite requirements.txt with the now-working versions.
"""

from __future__ import annotations

import importlib
import importlib.metadata as md
import subprocess
import sys
import traceback
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent
REQ_FILE = PROJECT_ROOT / "requirements.txt"

# Packages we know failed
FAILING = ["seqeval", "gradio", "fastapi", "wandb"]

PIP_BASE = [sys.executable, "-m", "pip", "install",
            "--user", "--break-system-packages", "--upgrade"]


# ---------------------------------------------------------------------------
# Step 1: diagnose
# ---------------------------------------------------------------------------

def diagnose(name: str) -> tuple[bool, str]:
    """Return (import_ok, traceback_or_version)."""
    # Force fresh import (in case a previous attempt cached a partial module)
    sys.modules.pop(name, None)
    try:
        mod = importlib.import_module(name)
        ver = getattr(mod, "__version__", "?")
        try:
            ver = md.version(name)
        except Exception:  # noqa: BLE001
            pass
        return True, ver
    except Exception:
        return False, traceback.format_exc()


def print_diagnosis() -> dict[str, tuple[bool, str]]:
    """Run diagnose() on each failing package and print the result."""
    print("=" * 72)
    print("STEP 1: Diagnose failing imports")
    print("=" * 72)
    results: dict[str, tuple[bool, str]] = {}
    for name in FAILING:
        ok, info = diagnose(name)
        results[name] = (ok, info)
        print(f"\n--- {name} ---")
        if ok:
            print(f"  Already importable. Version: {info}")
        else:
            # Print only the last few lines of the traceback (the cause)
            lines = info.strip().splitlines()
            tail = lines[-12:] if len(lines) > 12 else lines
            for ln in tail:
                print(f"  {ln}")
    return results


# ---------------------------------------------------------------------------
# Step 2: fixes
# ---------------------------------------------------------------------------

def run_pip(args: list[str], label: str) -> bool:
    """Run a pip install command and stream its tail to stdout."""
    cmd = PIP_BASE + args
    print(f"\n>>> {label}")
    print(f"    {' '.join(cmd)}")
    proc = subprocess.run(cmd, capture_output=True, text=True, check=False)
    log = (proc.stdout or "") + (proc.stderr or "")
    tail = "\n".join(log.strip().splitlines()[-15:])
    print(tail)
    return proc.returncode == 0


def fix_seqeval() -> bool:
    """seqeval fails to build because setuptools_scm 10.x removed an internal
    module. Pin setuptools_scm<10 and reinstall.
    """
    print("\n" + "=" * 72)
    print("FIX: seqeval (pin setuptools_scm<10)")
    print("=" * 72)
    ok1 = run_pip(["setuptools_scm<10"], "downgrade setuptools_scm")
    ok2 = run_pip(["--no-cache-dir", "seqeval==1.2.2"], "install seqeval")
    return ok1 and ok2


def fix_starlette_stack() -> bool:
    """gradio 6.14 + fastapi 0.136 + starlette 1.0.0 ended up incompatible.
    Reinstall a known-good stack: starlette<1.0, fastapi<0.115, gradio 4.44.1.

    gradio 4.44.1 is the last well-tested 4.x; pinning it removes the bleeding
    edge surface while keeping all features we use (Blocks, Chatbot, etc.).
    """
    print("\n" + "=" * 72)
    print("FIX: gradio + fastapi (pin to stable stack)")
    print("=" * 72)
    # First uninstall the bad versions to avoid resolver footguns
    uninstall = [sys.executable, "-m", "pip", "uninstall", "-y",
                 "--break-system-packages",
                 "starlette", "fastapi", "gradio", "gradio-client",
                 "hf-gradio", "safehttpx"]
    print(f"\n>>> uninstall conflicting packages")
    print(f"    {' '.join(uninstall)}")
    proc = subprocess.run(uninstall, capture_output=True, text=True, check=False)
    print("\n".join((proc.stdout or "").strip().splitlines()[-10:]))

    # Then reinstall a coherent set
    pkgs = [
        "starlette<1.0",
        "fastapi>=0.110,<0.115",
        "gradio==4.44.1",
    ]
    return run_pip(pkgs, "install pinned starlette/fastapi/gradio")


def fix_wandb() -> bool:
    """wandb 0.26 fails to import (likely protobuf 7 vs <5 mismatch).
    Pin protobuf<6 and reinstall wandb. Optional — failure is non-blocking.
    """
    print("\n" + "=" * 72)
    print("FIX: wandb (optional — pin protobuf<6)")
    print("=" * 72)
    ok1 = run_pip(["protobuf<6"], "downgrade protobuf")
    ok2 = run_pip(["wandb>=0.16,<0.20"], "install older wandb")
    return ok1 and ok2


# ---------------------------------------------------------------------------
# Step 3: verify + rewrite requirements
# ---------------------------------------------------------------------------

def verify_all() -> dict[str, tuple[bool, str]]:
    """Re-run diagnose() after fixes."""
    print("\n" + "=" * 72)
    print("STEP 3: Re-verify imports")
    print("=" * 72)
    results: dict[str, tuple[bool, str]] = {}
    for name in FAILING:
        ok, info = diagnose(name)
        results[name] = (ok, info)
        if ok:
            print(f"  ✓ {name:<20s} {info}")
        else:
            print(f"  ✗ {name:<20s} STILL FAILING")
    return results


def update_requirements(post: dict[str, tuple[bool, str]]) -> None:
    """Append a 'fixes' section to requirements.txt with the now-working pins."""
    if not REQ_FILE.exists():
        print("\n(no requirements.txt to update)")
        return
    extra = ["", "# --- post-fix pins (added by setup/fix_install.py) ---"]
    for name, (ok, info) in post.items():
        if ok:
            extra.append(f"{name}=={info}")
        else:
            extra.append(f"# {name}  (STILL FAILING: {info.splitlines()[-1] if info else 'unknown'})")
    REQ_FILE.write_text(REQ_FILE.read_text() + "\n".join(extra) + "\n")
    print(f"\nrequirements.txt updated -> {REQ_FILE}")


def main() -> int:
    """Orchestrate diagnose -> fix -> verify."""
    print("Multilingual Chatbot — Repair failed installs")
    print(f"Python: {sys.version.split()[0]}\n")

    pre = print_diagnosis()

    # Apply fixes only for the things actually broken
    if not pre["seqeval"][0]:
        fix_seqeval()
    if not pre["gradio"][0] or not pre["fastapi"][0]:
        fix_starlette_stack()
    if not pre["wandb"][0]:
        fix_wandb()

    post = verify_all()
    update_requirements(post)

    # Decide return code: only required ones (seqeval, gradio, fastapi) block
    blocking = [n for n in ("seqeval", "gradio", "fastapi") if not post[n][0]]
    if blocking:
        print(f"\nStill broken (BLOCKING): {blocking}")
        print("Paste the diagnostic output above so we can refine the fix.")
        return 1
    print("\nAll required packages now import cleanly. Safe to continue to Phase 2.")
    if not post["wandb"][0]:
        print("(wandb is still broken but is optional — we'll proceed without it.)")
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print("\nAborted by user.")
        sys.exit(130)