backend/app/studio/model_manager.py

"""
Studio Model Manager (essay-to-video pipeline, Batch 3).

Makes "download it from Civitai" an actual capability instead of the manual
wget + hand-edit process MODELS_README.md used to describe. Wraps the
official Civitai REST API (civitai.com - see sourcing rule 1) and Hugging
Face Hub file downloads, verifies hashes, resolves licenses against an
allowlist, and registers installs into a JSON registry that
providers.available_image_models()/available_video_models() already read -
no hand-edits, no code changes per model.

Three sourcing rules, from the design doc
(docs/design/essay-to-video/BATCH-3-model-manager.md):

1. civitai.com, never civitai.red. This pipeline is 100% SFW; the base URL
   is a constant, not a config knob.
2. Hugging Face first for base checkpoints, Civitai for style layers
   (LoRAs/finetunes). MODEL_SOURCE_PREFERENCE=huggingface_first records the
   policy; the two download paths exist either way.
3. The essay pipeline's picker never surfaces the mature-lane models
   (essay_pipeline_models() filters them; HomePilot's general pickers are
   untouched).

License rule: register() refuses any model without a resolved license row.
That's the guard against pulling something not cleared for a monetized
channel - a hard stop, not a warning.

ADDITIVE ONLY - reuses backend/app/civitai.py (CivitaiClient) and
backend/app/providers.py (models path + registry read); modifies neither.
"""
from __future__ import annotations

import hashlib
import json
import os
import re
import threading
import time
import uuid
from pathlib import Path
from typing import Any, Dict, List, Literal, Optional

from pydantic import BaseModel, Field

ModelType = Literal["image", "video", "lora"]


# ============================================================================
# License allowlist
# ============================================================================

class LicenseRow(BaseModel):
    """A resolved license decision for one model family."""
    key: str                     # substring matched against model/repo ids
    license: str
    commercial_ok: bool
    notes: str = ""


# Mid-2026 state, per the design doc's license table. Order matters: first
# substring match wins. A model that matches no row cannot be registered.
LICENSE_ALLOWLIST: List[LicenseRow] = [
    LicenseRow(key="flux.2-klein", license="Apache-2.0", commercial_ok=True,
               notes="Default image model for the essay pipeline."),
    LicenseRow(key="flux2-klein", license="Apache-2.0", commercial_ok=True),
    LicenseRow(key="flux.2-dev", license="BFL non-commercial", commercial_ok=False,
               notes="Paid license required for commercial use - manual hero shots only."),
    LicenseRow(key="flux2-dev", license="BFL non-commercial", commercial_ok=False),
    LicenseRow(key="stable-diffusion-3.5", license="Stability Community License",
               commercial_ok=True, notes="Free under $1M annual revenue."),
    LicenseRow(key="sd3.5", license="Stability Community License", commercial_ok=True),
    LicenseRow(key="wan2.2", license="Apache-2.0", commercial_ok=True),
    LicenseRow(key="wan-2.2", license="Apache-2.0", commercial_ok=True),
    LicenseRow(key="mochi", license="Apache-2.0", commercial_ok=True),
    LicenseRow(key="ltx-video", license="LTX community license", commercial_ok=True,
               notes="Verify current terms for monetized channels."),
    LicenseRow(key="hunyuanvideo", license="Tencent community license", commercial_ok=True,
               notes="Verify current terms for monetized channels."),
    LicenseRow(key="qwen-image", license="UNRESOLVED", commercial_ok=False,
               notes="Design doc: verify at integration time. Blocked until resolved."),
    # Rows below cover the essay-video bundle components (bundles.py)
    LicenseRow(key="sd_xl_base", license="CreativeML Open RAIL++-M", commercial_ok=True),
    LicenseRow(key="sdxl", license="CreativeML Open RAIL++-M", commercial_ok=True),
    LicenseRow(key="svd", license="Stability AI Community License", commercial_ok=True,
               notes="Free under $1M annual revenue."),
    LicenseRow(key="flux1-schnell", license="Apache-2.0", commercial_ok=True),
    LicenseRow(key="flux-schnell", license="Apache-2.0", commercial_ok=True),
    LicenseRow(key="hunyuan-video", license="Tencent community license", commercial_ok=True,
               notes="Verify current terms for monetized channels."),
    LicenseRow(key="ltx", license="LTX community license", commercial_ok=True,
               notes="Verify current terms for monetized channels."),
]

# Model IDs (and matching substrings) that belong to HomePilot's mature
# lane. The essay pipeline filters these out entirely - not just by default.
MATURE_LANE_PATTERNS = ("pony", "uncensored", "nsfw")


def resolve_license(model_id: str) -> Optional[LicenseRow]:
    """First allowlist row whose key appears in the model id (case-insensitive)."""
    needle = (model_id or "").lower()
    for row in LICENSE_ALLOWLIST:
        if row.key in needle:
            return row
    return None


# ============================================================================
# Data shapes
# ============================================================================

class CivitaiModelSummary(BaseModel):
    model_id: int
    version_id: Optional[int] = None
    name: str
    type: str = ""
    creator: str = ""
    download_count: int = 0
    allow_commercial_use: str = ""       # Civitai's own field, informative only
    nsfw: bool = False


class InstalledModel(BaseModel):
    id: str                              # picker id, e.g. "flux2-klein"
    model_type: ModelType
    filename: str
    path: str
    sha256: str = ""
    source: Literal["huggingface", "civitai"]
    source_ref: str = ""                 # repo_id or civitai version id
    license: str = ""
    commercial_ok: bool = False
    workflow: str = ""                   # ComfyUI workflow this maps to
    installed_at: float = Field(default_factory=time.time)


# ============================================================================
# Registry (read by providers.available_*_models(), written here)
# ============================================================================

def registry_path() -> Path:
    p = os.getenv("MODEL_REGISTRY_PATH", "").strip()
    if p:
        return Path(p)
    return Path(__file__).resolve().parents[2] / "data" / "installed_models.json"


def read_registry() -> List[Dict[str, Any]]:
    try:
        with open(registry_path(), "r", encoding="utf-8") as fh:
            data = json.load(fh)
        return data if isinstance(data, list) else []
    except Exception:
        return []


def _write_registry(entries: List[Dict[str, Any]]) -> None:
    path = registry_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp")
    with open(tmp, "w", encoding="utf-8") as fh:
        json.dump(entries, fh, indent=2)
    tmp.replace(path)


# ============================================================================
# The manager
# ============================================================================

class ModelManager:
    # Rule 1: the PG front door, hard-coded. Never civitai.red.
    CIVITAI_BASE = "https://civitai.com/api/v1"

    def __init__(self, api_key: Optional[str] = None, models_dir: Optional[Path] = None):
        self.api_key = (api_key or os.getenv("CIVITAI_API_KEY", "")).strip() or None
        if models_dir is None:
            from ..providers import get_comfy_models_path
            models_dir = get_comfy_models_path()
        self.models_dir = Path(models_dir)
        self.source_preference = os.getenv(
            "MODEL_SOURCE_PREFERENCE", "huggingface_first").strip()

    # ── Search ──────────────────────────────────────────────────────────────

    async def search_civitai(self, query: str, model_type: str = "image",
                             sfw_only: bool = True, limit: int = 20
                             ) -> List[CivitaiModelSummary]:
        """SFW-only search against civitai.com. sfw_only=False is accepted for
        signature compatibility but ignored - this pipeline never requests
        NSFW results (rule 1/3)."""
        from ..civitai import CivitaiClient, CivitaiSearchQuery

        client = CivitaiClient(api_key=self.api_key)
        raw = await client.search_models(CivitaiSearchQuery(
            query=query, model_type=model_type, limit=limit, nsfw=False))

        out: List[CivitaiModelSummary] = []
        for item in raw.get("items", []):
            if item.get("nsfw"):
                continue  # belt and braces on top of the SFW query
            versions = item.get("modelVersions") or []
            out.append(CivitaiModelSummary(
                model_id=item.get("id", 0),
                version_id=(versions[0].get("id") if versions else None),
                name=item.get("name", ""),
                type=item.get("type", ""),
                creator=(item.get("creator") or {}).get("username", ""),
                download_count=(item.get("stats") or {}).get("downloadCount", 0),
                allow_commercial_use=str(item.get("allowCommercialUse", "")),
                nsfw=bool(item.get("nsfw")),
            ))
        return out

    # ── Downloads ───────────────────────────────────────────────────────────

    def _download_stream(self, url: str, dest: Path,
                         headers: Optional[Dict[str, str]] = None) -> str:
        """Stream a file to dest, returning its sha256."""
        import httpx

        dest.parent.mkdir(parents=True, exist_ok=True)
        digest = hashlib.sha256()
        tmp = dest.with_suffix(dest.suffix + ".part")
        with httpx.stream("GET", url, headers=headers or {}, timeout=None,
                          follow_redirects=True) as resp:
            resp.raise_for_status()
            with open(tmp, "wb") as fh:
                for chunk in resp.iter_bytes(chunk_size=1 << 20):
                    fh.write(chunk)
                    digest.update(chunk)
        tmp.replace(dest)
        return digest.hexdigest()

    def download_from_huggingface(self, repo_id: str, filename: str,
                                  model_type: ModelType = "image",
                                  subdir: str = "checkpoints",
                                  revision: str = "main") -> InstalledModel:
        """Rule 2: base checkpoints come from official Hugging Face repos."""
        license_row = resolve_license(f"{repo_id}/{filename}")
        if license_row is None:
            raise PermissionError(
                f"No resolved license row for '{repo_id}/{filename}'. "
                "Add it to LICENSE_ALLOWLIST before installing (design §7.3).")

        url = f"https://huggingface.co/{repo_id}/resolve/{revision}/{filename}"
        headers = {}
        hf_token = os.getenv("HF_TOKEN", "").strip()
        if hf_token:
            headers["Authorization"] = f"Bearer {hf_token}"

        dest = self.models_dir / subdir / filename
        sha = self._download_stream(url, dest, headers)

        return InstalledModel(
            id=_picker_id(filename), model_type=model_type,
            filename=filename, path=str(dest), sha256=sha,
            source="huggingface", source_ref=repo_id,
            license=license_row.license, commercial_ok=license_row.commercial_ok,
        )

    async def download_by_version_id(self, version_id: int,
                                     model_type: ModelType = "lora",
                                     subdir: str = "loras") -> InstalledModel:
        """Rule 2: Civitai is the style-layer source (LoRAs/finetunes)."""
        import httpx

        headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
        async with httpx.AsyncClient(timeout=30.0) as client:
            r = await client.get(f"{self.CIVITAI_BASE}/model-versions/{version_id}",
                                 headers=headers)
            r.raise_for_status()
            meta = r.json()

        model_meta = meta.get("model") or {}
        if model_meta.get("nsfw"):
            raise PermissionError("NSFW model refused: the essay pipeline is SFW-only.")

        name = f"{model_meta.get('name', 'model')} @ {meta.get('name', version_id)}"
        license_row = resolve_license(name) or self._license_from_civitai(model_meta)
        if license_row is None:
            raise PermissionError(
                f"No resolvable license for Civitai version {version_id} ('{name}'). "
                "Refusing to install (design §7.3).")

        files = meta.get("files") or []
        primary = next((f for f in files if f.get("primary")), files[0] if files else None)
        if not primary:
            raise ValueError(f"Civitai version {version_id} has no downloadable files")

        filename = primary.get("name") or f"civitai-{version_id}.safetensors"
        url = primary.get("downloadUrl") or f"{self.CIVITAI_BASE}/download/models/{version_id}"
        dest = self.models_dir / subdir / filename
        sha = self._download_stream(url, dest, headers)

        declared = ((primary.get("hashes") or {}).get("SHA256") or "").lower()
        if declared and declared != sha:
            dest.unlink(missing_ok=True)
            raise ValueError(f"SHA256 mismatch for {filename}: "
                             f"declared {declared[:12]}…, got {sha[:12]}…")

        return InstalledModel(
            id=_picker_id(filename), model_type=model_type,
            filename=filename, path=str(dest), sha256=sha,
            source="civitai", source_ref=str(version_id),
            license=license_row.license, commercial_ok=license_row.commercial_ok,
        )

    @staticmethod
    def _license_from_civitai(model_meta: Dict[str, Any]) -> Optional[LicenseRow]:
        """Derive a license row from Civitai's allowCommercialUse when the
        model isn't in the static allowlist. 'None'/missing -> unresolvable."""
        allow = model_meta.get("allowCommercialUse")
        values = {str(v).lower() for v in (allow if isinstance(allow, list) else [allow])}
        if values & {"image", "rentcivit", "rent", "sell"}:
            return LicenseRow(key="civitai-declared", license="Civitai: commercial use declared",
                              commercial_ok=True,
                              notes=f"allowCommercialUse={sorted(values)}")
        return None

    # ── Registration ────────────────────────────────────────────────────────

    def register(self, model: InstalledModel, workflow: str = "") -> None:
        """Make an installed model selectable. Refuses unlicensed models
        (defense in depth - the download paths already check)."""
        if not model.license or model.license == "UNRESOLVED":
            raise PermissionError(
                f"Refusing to register '{model.id}': no resolved license row.")
        if model.workflow == "" and workflow:
            model.workflow = workflow

        entries = read_registry()
        entries = [e for e in entries if e.get("id") != model.id]
        entries.append(json.loads(model.model_dump_json()))
        _write_registry(entries)

    def installed(self) -> List[InstalledModel]:
        return [InstalledModel(**e) for e in read_registry()]


def _picker_id(filename: str) -> str:
    """checkpoint filename -> stable picker id (lowercase, no extension)."""
    stem = Path(filename).stem.lower()
    return re.sub(r"[^a-z0-9.]+", "-", stem).strip("-")


# ============================================================================
# Essay-pipeline model filtering (rule 3)
# ============================================================================

def essay_pipeline_models(model_type: str = "image") -> List[str]:
    """The model list the essay pipeline's picker shows: the general list
    minus the mature lane, entirely - not just defaulted away."""
    from ..providers import available_image_models, available_video_models

    models = available_image_models() if model_type == "image" else available_video_models()
    return [m for m in models
            if not any(pat in m.lower() for pat in MATURE_LANE_PATTERNS)]


# ============================================================================
# Background install jobs (downloads are long; endpoints return 202 + poll)
# ============================================================================

_JOBS: Dict[str, Dict[str, Any]] = {}
_JOBS_LOCK = threading.Lock()


def submit_install(kind: str, **kwargs: Any) -> Dict[str, Any]:
    """kind: 'huggingface' (repo_id, filename, model_type, subdir) or
    'civitai' (version_id, model_type, subdir). Returns the job descriptor."""
    job_id = str(uuid.uuid4())
    job = {"id": job_id, "kind": kind, "status": "running",
           "params": kwargs, "error": "", "model": None,
           "started_at": time.time()}
    with _JOBS_LOCK:
        _JOBS[job_id] = job

    def _run() -> None:
        try:
            mgr = ModelManager()
            if kind == "huggingface":
                model = mgr.download_from_huggingface(
                    repo_id=kwargs["repo_id"], filename=kwargs["filename"],
                    model_type=kwargs.get("model_type", "image"),
                    subdir=kwargs.get("subdir", "checkpoints"))
            elif kind == "civitai":
                import asyncio
                model = asyncio.run(mgr.download_by_version_id(
                    version_id=int(kwargs["version_id"]),
                    model_type=kwargs.get("model_type", "lora"),
                    subdir=kwargs.get("subdir", "loras")))
            else:
                raise ValueError(f"Unknown install kind: {kind}")
            mgr.register(model, workflow=kwargs.get("workflow", ""))
            with _JOBS_LOCK:
                job["status"] = "done"
                job["model"] = json.loads(model.model_dump_json())
        except Exception as e:
            with _JOBS_LOCK:
                job["status"] = "error"
                job["error"] = str(e)

    threading.Thread(target=_run, daemon=True).start()
    return {"id": job_id, "status": "running"}


def get_install_job(job_id: str) -> Optional[Dict[str, Any]]:
    with _JOBS_LOCK:
        job = _JOBS.get(job_id)
        return dict(job) if job else None