"""
CorridorKey GPU Server — setup and launch as a subprocess.

Uses a dedicated venv at corridor-key/.venv/ so dependencies persist
even when sorcgcs-server itself runs inside uvx's ephemeral environment.
"""

import builtins
import os
import subprocess
import sys
import shutil
import time
import platform
import venv

# Force all prints to flush immediately (Windows buffers stdout when not a tty)
_builtin_print = builtins.print
def print(*args, **kwargs):
    kwargs.setdefault("flush", True)
    _builtin_print(*args, **kwargs)

CK_PORT = 8100
CK_IMG_SIZE = 512
HF_CHECKPOINT_URL = "https://huggingface.co/nikopueringer/CorridorKey_v1.0/resolve/main/CorridorKey_v1.0.pth"
GH_REPO_URL = "https://github.com/nikopueringer/CorridorKey.git"
GH_REPO_COMMIT = "e6236e1a2a8f495049eb92f24c623a7ebd6a9a53"
GH_ZIP_URL = "https://github.com/nikopueringer/CorridorKey/archive/e6236e1a2a8f495049eb92f24c623a7ebd6a9a53.zip"


def get_ck_dir():
    home = os.path.expanduser("~")
    return os.path.join(home, ".sorcgcs", "corridor-key")


def _venv_python():
    """Return the path to the Python inside corridor-key/.venv/."""
    ck_dir = get_ck_dir()
    if platform.system() == "Windows":
        return os.path.join(ck_dir, ".venv", "Scripts", "python.exe")
    return os.path.join(ck_dir, ".venv", "bin", "python")


def _check_file(path):
    return os.path.isfile(path)


def _status():
    """Return a dict describing what's already set up."""
    ck_dir = get_ck_dir()
    return {
        "venv": _check_file(_venv_python()),
        "repo": _check_file(os.path.join(ck_dir, "CorridorKey", "CorridorKeyModule", "__init__.py")),
        "checkpoint": _check_file(os.path.join(ck_dir, "CorridorKey", "CorridorKeyModule", "checkpoints", "CorridorKey.pth")),
        "server_py": _check_file(os.path.join(ck_dir, "server.py")),
        "torch": False,  # checked lazily
    }


def _torch_installed():
    """Check if torch is importable in the CK venv (and the venv is healthy)."""
    py = _venv_python()
    cfg = os.path.join(get_ck_dir(), ".venv", "pyvenv.cfg")
    if not _check_file(py) or not _check_file(cfg):
        return False
    r = subprocess.run([py, "-c", "import torch"], capture_output=True)
    return r.returncode == 0


def is_fully_ready():
    s = _status()
    if not (s["venv"] and s["repo"] and s["checkpoint"] and s["server_py"]):
        return False
    return _torch_installed()


# ── Step helpers ────────────────────────────────────────────────────────────

def _create_venv():
    ck_dir = get_ck_dir()
    venv_dir = os.path.join(ck_dir, ".venv")
    py = _venv_python()
    cfg = os.path.join(venv_dir, "pyvenv.cfg")

    # Check if venv is healthy (both python.exe AND pyvenv.cfg must exist)
    if _check_file(py) and _check_file(cfg):
        return True

    # Broken or missing venv — wipe and recreate
    if os.path.isdir(venv_dir):
        print("         Removing broken venv...")
        import shutil as _shutil
        _shutil.rmtree(venv_dir, ignore_errors=True)
        time.sleep(1)

    print("         Creating Python virtual environment...")
    os.makedirs(ck_dir, exist_ok=True)

    # Use subprocess with --clear for maximum reliability on Windows
    r = subprocess.run(
        [sys.executable, "-m", "venv", "--clear", venv_dir],
        capture_output=True, text=True,
    )
    if r.returncode != 0:
        print(f"         venv command failed: {r.stderr[:300]}")
        # Fallback to venv module
        try:
            venv.create(venv_dir, with_pip=True, clear=True)
        except Exception as e:
            print(f"         FAILED: Could not create venv: {e}")
            return False

    if not _check_file(cfg):
        print(f"         FAILED: pyvenv.cfg not created at {cfg}")
        return False
    if not _check_file(py):
        print(f"         FAILED: Python not found at {py}")
        return False

    # Upgrade pip
    subprocess.run([py, "-m", "pip", "install", "--upgrade", "pip"], capture_output=True)
    print("         venv OK.")
    return True


def _clone_repo():
    ck_dir = get_ck_dir()
    repo_dir = os.path.join(ck_dir, "CorridorKey")
    init_py = os.path.join(repo_dir, "CorridorKeyModule", "__init__.py")

    # If repo exists, verify it's on the correct commit
    if _check_file(init_py) and shutil.which("git") and os.path.isdir(os.path.join(repo_dir, ".git")):
        r = subprocess.run(["git", "rev-parse", "HEAD"], cwd=repo_dir, capture_output=True, text=True)
        if r.returncode == 0 and r.stdout.strip() == GH_REPO_COMMIT:
            return True
        print("         CorridorKey repo is wrong version, re-cloning...")
        shutil.rmtree(repo_dir, ignore_errors=True)
        time.sleep(1)
    elif _check_file(init_py):
        return True

    os.makedirs(ck_dir, exist_ok=True)

    if shutil.which("git"):
        print("         Cloning from GitHub...")
        r = subprocess.run(
            ["git", "clone", GH_REPO_URL, repo_dir],
            capture_output=True, text=True,
        )
        if r.returncode == 0:
            # Pin to known-good commit
            subprocess.run(
                ["git", "checkout", GH_REPO_COMMIT],
                cwd=repo_dir, capture_output=True, text=True,
            )
            if _check_file(os.path.join(repo_dir, "CorridorKeyModule", "__init__.py")):
                return True
        print(f"         Git clone failed: {r.stderr[:200]}")

    # Fallback: ZIP download
    print("         Downloading as ZIP...")
    zip_path = os.path.join(ck_dir, "CorridorKey.zip")
    downloaded = False

    if shutil.which("curl"):
        r = subprocess.run(["curl", "-sL", "-o", zip_path, GH_ZIP_URL], capture_output=True)
        downloaded = r.returncode == 0 and _check_file(zip_path)

    if not downloaded:
        try:
            import urllib.request
            urllib.request.urlretrieve(GH_ZIP_URL, zip_path)
            downloaded = _check_file(zip_path)
        except Exception as e:
            print(f"         Download failed: {e}")
            return False

    if not downloaded:
        print("         FAILED: Could not download CorridorKey.")
        return False

    import zipfile
    with zipfile.ZipFile(zip_path, "r") as z:
        z.extractall(ck_dir)
    # ZIP extracts as CorridorKey-<hash> or CorridorKey-main
    extracted = os.path.join(ck_dir, "CorridorKey-" + GH_REPO_COMMIT)
    if not os.path.isdir(extracted):
        extracted = os.path.join(ck_dir, "CorridorKey-main")
    if os.path.isdir(extracted):
        if os.path.isdir(repo_dir):
            shutil.rmtree(repo_dir)
        os.rename(extracted, repo_dir)
    if _check_file(zip_path):
        os.remove(zip_path)
    return _check_file(init_py)


def _download_checkpoint():
    ck_dir = get_ck_dir()
    ckpt_dir = os.path.join(ck_dir, "CorridorKey", "CorridorKeyModule", "checkpoints")
    ckpt_file = os.path.join(ckpt_dir, "CorridorKey.pth")
    if _check_file(ckpt_file) and os.path.getsize(ckpt_file) > 1_000_000:
        return True

    os.makedirs(ckpt_dir, exist_ok=True)
    print("         Downloading from HuggingFace (~300 MB)...")

    if shutil.which("curl"):
        r = subprocess.run(["curl", "-L", "--progress-bar", "-o", ckpt_file, HF_CHECKPOINT_URL])
        if r.returncode == 0 and _check_file(ckpt_file) and os.path.getsize(ckpt_file) > 1_000_000:
            return True

    try:
        import urllib.request
        print("         (using Python urllib — this may take a while)")
        urllib.request.urlretrieve(HF_CHECKPOINT_URL, ckpt_file)
        if _check_file(ckpt_file) and os.path.getsize(ckpt_file) > 1_000_000:
            return True
    except Exception as e:
        print(f"         Download error: {e}")

    if _check_file(ckpt_file):
        os.remove(ckpt_file)
    return False


def _install_deps():
    py = _venv_python()
    if not _check_file(py):
        return False

    # Check if torch already works in the venv
    r = subprocess.run([py, "-c", "import torch; print(torch.__version__)"], capture_output=True, text=True)
    torch_ok = r.returncode == 0

    if not torch_ok:
        print("         Installing PyTorch (large download, may take several minutes)...")
        print("         You will see pip output below:")
        print("")
        is_mac = platform.system() == "Darwin"

        if is_mac:
            subprocess.run([py, "-m", "pip", "install", "torch", "torchvision"])
        else:
            installed = False
            for cuda_ver in ["cu124", "cu121"]:
                print(f"\n         --- Trying PyTorch + {cuda_ver} ---")
                r = subprocess.run(
                    [py, "-m", "pip", "install", "torch", "torchvision",
                     "--index-url", f"https://download.pytorch.org/whl/{cuda_ver}"],
                )
                if r.returncode == 0:
                    check = subprocess.run(
                        [py, "-c", "import torch; assert torch.cuda.is_available()"],
                        capture_output=True,
                    )
                    if check.returncode == 0:
                        print(f"         PyTorch + {cuda_ver} installed — CUDA available!")
                        installed = True
                        break
                    else:
                        print(f"         {cuda_ver} installed but CUDA not detected, trying next...")
                        subprocess.run([py, "-m", "pip", "uninstall", "torch", "torchvision", "-y"],
                                       capture_output=True)

            if not installed:
                print("\n         --- Installing CPU-only PyTorch as fallback ---")
                subprocess.run([py, "-m", "pip", "install", "torch", "torchvision"])

        print("")

    # Install remaining deps (these are small, capture output)
    print("         Installing FastAPI, uvicorn, opencv, etc...")
    subprocess.run(
        [py, "-m", "pip", "install", "timm", "fastapi", "uvicorn[standard]",
         "pydantic", "opencv-python", "numpy", "pillow"],
        capture_output=True,
    )

    # Verify torch actually installed
    r = subprocess.run(
        [py, "-c",
         "import torch; cuda=torch.cuda.is_available(); "
         "mps=hasattr(torch.backends,'mps') and torch.backends.mps.is_available(); "
         "print(f'{torch.__version__}|{cuda}|{mps}')"],
        capture_output=True, text=True,
    )
    if r.returncode == 0:
        parts = r.stdout.strip().split("|")
        print(f"         PyTorch {parts[0]} — CUDA: {parts[1]}, MPS: {parts[2]}")
        return {"version": parts[0], "cuda": parts[1] == "True", "mps": parts[2] == "True"}

    print("         ERROR: PyTorch failed to install.")
    return False


def _verify_ck_repo_version():
    """If the CorridorKey repo exists but is on the wrong commit, delete it so setup re-clones."""
    ck_dir = get_ck_dir()
    repo_dir = os.path.join(ck_dir, "CorridorKey")
    git_dir = os.path.join(repo_dir, ".git")
    if not os.path.isdir(git_dir):
        return
    if not shutil.which("git"):
        return
    r = subprocess.run(["git", "rev-parse", "HEAD"], cwd=repo_dir, capture_output=True, text=True)
    if r.returncode == 0 and r.stdout.strip() == GH_REPO_COMMIT:
        return
    print("  CorridorKey repo is outdated, removing for re-clone...")
    shutil.rmtree(repo_dir, ignore_errors=True)
    time.sleep(1)


def _write_server_py():
    ck_dir = get_ck_dir()
    os.makedirs(ck_dir, exist_ok=True)
    path = os.path.join(ck_dir, "server.py")
    with open(path, "w", encoding="utf-8") as f:
        f.write(CK_SERVER_SOURCE)
    return _check_file(path)


# ── Main entry points ──────────────────────────────────────────────────────

def setup_corridorkey():
    """Full first-time setup. Returns a result dict."""
    result = {"ok": False, "error": None, "torch_info": None}

    steps = [
        ("Creating virtual environment", _create_venv),
        ("Downloading CorridorKey repo", _clone_repo),
        ("Downloading model checkpoint", _download_checkpoint),
        ("Installing Python dependencies", _install_deps),
        ("Writing server script", _write_server_py),
    ]

    print("")
    print("  CorridorKey first-time setup")
    print("  (GPU green screen removal for AutoSprite)")
    print("  Needs: ~2 GB disk, GPU recommended (NVIDIA CUDA / Apple MPS)")
    print("")

    for i, (label, fn) in enumerate(steps, 1):
        print(f"  [{i}/{len(steps)}] {label}...")
        try:
            ret = fn()
        except Exception as e:
            result["error"] = f"Step {i} ({label}) failed: {e}"
            print(f"         FAILED: {e}")
            return result
        if ret is False:
            result["error"] = f"Step {i} ({label}) failed."
            return result
        if isinstance(ret, dict) and "version" in ret:
            result["torch_info"] = ret

    print("")
    result["ok"] = True
    return result


def launch_corridorkey():
    """
    Set up (if needed) and launch the CK server.
    Returns a dict: {"proc": Popen|None, "status": str, "error": str|None, "gpu": str|None}
    """
    info = {"proc": None, "status": "skipped", "error": None, "gpu": None}

    # Always write latest server.py
    _write_server_py()

    # Check if CorridorKey repo is on the correct commit — force re-clone if wrong
    _verify_ck_repo_version()

    if not is_fully_ready():
        setup_result = setup_corridorkey()
        if not setup_result["ok"]:
            info["status"] = "setup_failed"
            info["error"] = setup_result.get("error", "Unknown setup error")
            return info

    # Kill any existing process on the CK port
    _kill_port(CK_PORT)

    py = _venv_python()
    ck_dir = get_ck_dir()
    server_py = os.path.join(ck_dir, "server.py")

    # Write CK output to a log file so we can read it if something goes wrong
    log_path = os.path.join(ck_dir, "server.log")
    log_file = open(log_path, "w", encoding="utf-8")

    proc = subprocess.Popen(
        [py, server_py, "--img-size", str(CK_IMG_SIZE), "--device", "auto", "--port", str(CK_PORT)],
        cwd=ck_dir,
        stdout=log_file,
        stderr=subprocess.STDOUT,
    )
    info["proc"] = proc

    # Poll for health — 5 min timeout for first run (torch dynamo compilation is slow)
    print(f"  Waiting for CorridorKey model to load (port {CK_PORT})...")
    print(f"  (First run may take several minutes for GPU compilation)")
    print(f"  Log: {log_path}")
    verified = False
    last_log_pos = 0
    for attempt in range(150):  # 150 * 2s = 5 minutes
        time.sleep(2)

        # Print only meaningful lines from CK server log (skip torch dynamo spam)
        try:
            log_file.flush()
            with open(log_path, "r", encoding="utf-8", errors="replace") as lf:
                lf.seek(last_log_pos)
                new_content = lf.read()
                last_log_pos = lf.tell()
            for line in new_content.splitlines():
                stripped = line.strip()
                if stripped and "[CK Server]" in stripped:
                    print(f"  {stripped}")
                elif stripped and ("ERROR" in stripped or "Error" in stripped) and "dynamo" not in stripped.lower():
                    print(f"  {stripped}")
        except Exception:
            pass

        if proc.poll() is not None:
            info["status"] = "crashed"
            try:
                with open(log_path, "r", encoding="utf-8", errors="replace") as lf:
                    info["error"] = f"Process exited with code {proc.returncode}.\n{lf.read()[-800:]}"
            except Exception:
                info["error"] = f"Process exited with code {proc.returncode}."
            log_file.close()
            return info

        try:
            import urllib.request
            import json
            resp = urllib.request.urlopen(f"http://localhost:{CK_PORT}/health", timeout=2)
            if resp.status == 200:
                data = json.loads(resp.read().decode())
                info["gpu"] = data.get("gpu_name") or data.get("device", "unknown")
                verified = True
                break
        except Exception:
            pass

    if verified:
        info["status"] = "running"
    else:
        info["status"] = "timeout"
        try:
            with open(log_path, "r", encoding="utf-8", errors="replace") as lf:
                log_tail = lf.read()[-800:]
            info["error"] = f"/health never returned 200 within 5 minutes.\nLog output:\n{log_tail}"
        except Exception:
            info["error"] = "/health never returned 200 within 5 minutes."

    log_file.close()
    return info


def _kill_port(port):
    try:
        if platform.system() == "Windows":
            r = subprocess.run(["netstat", "-aon"], capture_output=True, text=True)
            for line in r.stdout.splitlines():
                if f":{port} " in line and "LISTENING" in line:
                    pid = line.split()[-1]
                    subprocess.run(["taskkill", "/PID", pid, "/F"], capture_output=True)
                    time.sleep(1)
        else:
            r = subprocess.run(["lsof", f"-ti:{port}"], capture_output=True, text=True)
            for pid in r.stdout.strip().split("\n"):
                if pid.strip():
                    subprocess.run(["kill", "-9", pid.strip()], capture_output=True)
            if r.stdout.strip():
                time.sleep(1)
    except Exception:
        pass


# ── Embedded CorridorKey server.py ──────────────────────────────────────────
CK_SERVER_SOURCE = r'''"""
CorridorKey FastAPI Server
==========================
Minimal GPU inference server that wraps the CorridorKey neural green screen keyer.
"""

import argparse
import base64
import io
import os
import sys
import time
from contextlib import asynccontextmanager

import torch._dynamo
torch._dynamo.config.suppress_errors = True

import cv2
import numpy as np
from PIL import Image

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
CK_REPO_DIR = os.path.join(SCRIPT_DIR, "CorridorKey")
if os.path.isdir(CK_REPO_DIR):
    sys.path.insert(0, CK_REPO_DIR)

parser = argparse.ArgumentParser(description="CorridorKey inference server")
parser.add_argument("--img-size", type=int, default=1024)
parser.add_argument("--device", type=str, default="auto")
parser.add_argument("--port", type=int, default=8100)
parser.add_argument("--host", type=str, default="0.0.0.0")
parser.add_argument("--checkpoint", type=str, default=None)
cli_args = parser.parse_args()

def resolve_device(requested):
    if requested != "auto":
        return requested
    try:
        import torch
        if torch.cuda.is_available():
            return "cuda"
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps"
    except ImportError:
        pass
    return "cpu"

def resolve_checkpoint(explicit):
    if explicit and os.path.isfile(explicit):
        return explicit
    candidates = [
        os.path.join(CK_REPO_DIR, "CorridorKeyModule", "checkpoints", "CorridorKey.pth"),
        os.path.join(CK_REPO_DIR, "CorridorKeyModule", "checkpoints", "CorridorKey_v1.0.pth"),
        os.path.join(SCRIPT_DIR, "CorridorKey.pth"),
    ]
    for path in candidates:
        if os.path.isfile(path):
            return path
    raise FileNotFoundError("CorridorKey checkpoint not found.")

engine = None
engine_device = None
engine_img_size = None

@asynccontextmanager
async def lifespan(app):
    global engine, engine_device, engine_img_size
    from CorridorKeyModule import CorridorKeyEngine
    device = resolve_device(cli_args.device)
    checkpoint = resolve_checkpoint(cli_args.checkpoint)
    img_size = cli_args.img_size
    print(f"[CK Server] Loading model: device={device}, img_size={img_size}")
    print(f"[CK Server] Checkpoint: {checkpoint}")
    engine = CorridorKeyEngine(checkpoint_path=checkpoint, device=device, img_size=img_size)
    engine_device = device
    engine_img_size = img_size
    print(f"[CK Server] Model loaded successfully on {device}")
    if device == "cuda":
        import torch
        allocated = torch.cuda.memory_allocated() / 1e9
        reserved = torch.cuda.memory_reserved() / 1e9
        print(f"[CK Server] VRAM: {allocated:.1f} GB allocated, {reserved:.1f} GB reserved")
    yield
    engine = None
    print("[CK Server] Shutdown complete")

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

app = FastAPI(title="CorridorKey Server", lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

class ProcessRequest(BaseModel):
    image: str
    alpha_hint: str
    resolution: int = 1024

def decode_base64_image(data):
    if "," in data: data = data.split(",", 1)[1]
    raw = base64.b64decode(data)
    return np.array(Image.open(io.BytesIO(raw)).convert("RGB"), dtype=np.float32) / 255.0

def decode_base64_mask(data):
    if "," in data: data = data.split(",", 1)[1]
    raw = base64.b64decode(data)
    return np.array(Image.open(io.BytesIO(raw)).convert("L"), dtype=np.float32) / 255.0

def encode_image_base64(arr, is_rgba=False):
    arr_uint8 = np.clip(arr * 255, 0, 255).astype(np.uint8)
    if is_rgba:
        pil_img = Image.fromarray(arr_uint8, "RGBA")
    elif arr_uint8.ndim == 2 or (arr_uint8.ndim == 3 and arr_uint8.shape[2] == 1):
        if arr_uint8.ndim == 3: arr_uint8 = arr_uint8[:, :, 0]
        pil_img = Image.fromarray(arr_uint8, "L")
    else:
        pil_img = Image.fromarray(arr_uint8, "RGB")
    buf = io.BytesIO()
    pil_img.save(buf, format="PNG")
    return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode('ascii')}"

@app.get("/health")
async def health():
    loaded = engine is not None
    info = {"status": "ok" if loaded else "loading", "device": engine_device, "img_size": engine_img_size, "model_loaded": loaded}
    if engine_device == "cuda":
        import torch
        info["gpu_name"] = torch.cuda.get_device_name(0)
        info["vram_allocated_gb"] = round(torch.cuda.memory_allocated() / 1e9, 2)
        info["vram_reserved_gb"] = round(torch.cuda.memory_reserved() / 1e9, 2)
        props = torch.cuda.get_device_properties(0)
        info["vram_total_gb"] = round(getattr(props, "total_memory", 0) / 1e9, 2)
    if not loaded:
        raise HTTPException(status_code=503, detail="Model still loading")
    return info

@app.post("/process")
async def process(req: ProcessRequest):
    if engine is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet")
    try:
        t0 = time.time()
        image = decode_base64_image(req.image)
        mask = decode_base64_mask(req.alpha_hint)
        h, w = image.shape[:2]
        print(f"[CK Server] Processing {w}x{h} frame at model resolution {engine_img_size}...")
        result = engine.process_frame(image=image, mask_linear=mask, refiner_scale=1.0, input_is_linear=False, fg_is_straight=True, despill_strength=1.0, auto_despeckle=True, despeckle_size=400)
        elapsed = time.time() - t0
        print(f"[CK Server] Done in {elapsed:.2f}s")
        fg_b64 = encode_image_base64(result["fg"])
        alpha_arr = result["alpha"]
        if alpha_arr.ndim == 3 and alpha_arr.shape[2] == 1: alpha_arr = alpha_arr[:, :, 0]
        alpha_b64 = encode_image_base64(alpha_arr)
        composite_b64 = encode_image_base64(result["comp"])
        rgba_arr = np.dstack([result["fg"], alpha_arr])
        processed_b64 = encode_image_base64(rgba_arr, is_rgba=True)
        return {"foreground": fg_b64, "alpha": alpha_b64, "composite": composite_b64, "processed": processed_b64, "processing_time_s": round(elapsed, 3)}
    except HTTPException:
        raise
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Processing failed: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    print(f"[CK Server] Starting on {cli_args.host}:{cli_args.port}")
    uvicorn.run(app, host=cli_args.host, port=cli_args.port)
'''
