"""Modal backend for FrogQuest's GPU work — the off-Space alternative to ZeroGPU.

Deploy once:  `modal deploy modal_app.py`
The HF Space (running on CPU-basic with FROGQUEST_BACKEND=modal) then invokes these via
`modal.Cls.from_name("frogquest", "LLM" | "Flux")().<method>.remote(...)` from llm.py / images.py.

Both backends share the SAME prompts, config, JSON extractor (gpu_shared) and JSON schemas
(schema) so they can never drift. Those two local files are mounted into the image with
`add_local_python_source`.

GPU config:
  - LLM (Nemotron Q8 ~4.3GB): a cheap L4 is plenty.
  - FLUX.2 klein (~23GB): set FROGQUEST_MODAL_FLUX_GPU at DEPLOY time (default "A10G", 24GB ->
    needs enable_model_cpu_offload(); use "L4" to save cost or "L40S"/"A100-40GB" if it OOMs).

NOTE: re-verify Modal API specifics against current docs (Modal 1.0+): App / Image.pip_install
extra_index_url / add_local_python_source / @app.cls / @modal.enter / @modal.method /
Cls.from_name / Volume.commit.
"""
import os

import modal

from gpu_shared import (
    GGUF_FILE,
    GGUF_REPO,
    GUIDANCE,
    INTENT_SYSTEM_PROMPT,
    MAX_SIDE,
    MODEL_ID,
    N_CTX,
    STEPS,
    SYSTEM_PROMPT,
    build_edit_prompt,
    build_initial_prompt,
    extract_json,
    preload_cuda_libs,
)
from schema import INTENT_SCHEMA, RESPONSE_SCHEMA

app = modal.App("frogquest")

# Persist the HF weight cache (FLUX ~23GB + GGUF ~4.3GB) across cold starts -> download once.
CACHE_DIR = "/cache"
vol = modal.Volume.from_name("frogquest-cache", create_if_missing=True)

# Prebuilt CUDA wheel for llama.cpp (same index the Space uses). cu124 wheels run on Modal GPUs.
_LLAMA_INDEX = "https://abetlen.github.io/llama-cpp-python/whl/cu124"

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "git+https://github.com/huggingface/diffusers.git",  # Flux2KleinPipeline (>=0.38)
        "transformers>=4.44.2",
        "accelerate",            # required by enable_model_cpu_offload()
        "torch",
        "pillow",
        "sentencepiece",         # FLUX.2 text-encoder tokenizer deps (belt-and-suspenders)
        "protobuf",
        "huggingface-hub",
        "hf-transfer",
    )
    .pip_install(
        "llama-cpp-python==0.3.28",
        extra_index_url=_LLAMA_INDEX,
        extra_options="--prefer-binary",
    )
    .env({"HF_HOME": CACHE_DIR, "HF_HUB_ENABLE_HF_TRANSFER": "1"})
    .add_local_python_source("schema", "gpu_shared")
)

# Configurable at deploy time (read on your machine when you run `modal deploy`).
FLUX_GPU = os.environ.get("FROGQUEST_MODAL_FLUX_GPU", "A10G")


@app.cls(gpu="L4", image=image, volumes={CACHE_DIR: vol}, scaledown_window=300)
class LLM:
    @modal.enter()
    def _load(self):
        vol.reload()  # pick up weights another container may have already cached
        import torch  # noqa: F401  (loads CUDA libs RTLD_GLOBAL)
        preload_cuda_libs()
        from llama_cpp import Llama

        self.llm = Llama.from_pretrained(
            repo_id=GGUF_REPO,
            filename=GGUF_FILE,
            n_gpu_layers=-1,
            n_ctx=N_CTX,  # Modal GPU has the VRAM for the full 128k
            verbose=False,
        )
        vol.commit()  # persist the freshly downloaded GGUF (no-op if already cached)

    @modal.method()
    def generate_quests(self, todos: str, theme: str) -> dict:
        system = SYSTEM_PROMPT.replace("{theme}", theme)
        user = f"Theme: {theme}\nMy to-do list / goals:\n{todos.strip()}"
        out = self.llm.create_chat_completion(
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
            response_format={"type": "json_object", "schema": RESPONSE_SCHEMA},
            temperature=0.0,
            max_tokens=4096,
        )
        return extract_json(out["choices"][0]["message"]["content"])

    @modal.method()
    def route_intent(self, message: str, context: str) -> dict:
        user = f"Context:\n{context.strip()}\n\nUser message:\n{message.strip()}"
        out = self.llm.create_chat_completion(
            messages=[
                {"role": "system", "content": INTENT_SYSTEM_PROMPT},
                {"role": "user", "content": user},
            ],
            response_format={"type": "json_object", "schema": INTENT_SCHEMA},
            temperature=0.0,
            max_tokens=256,
        )
        parsed = extract_json(out["choices"][0]["message"]["content"])
        if not isinstance(parsed, dict) or parsed.get("intent") not in (
            "forge", "add_tasks", "mark_done", "mark_couldnt", "unknown",
        ):
            return {"intent": "unknown"}
        return parsed


@app.cls(gpu=FLUX_GPU, image=image, volumes={CACHE_DIR: vol}, scaledown_window=300)
class Flux:
    @modal.enter()
    def _load(self):
        vol.reload()
        import torch
        from diffusers import Flux2KleinPipeline

        self._torch = torch
        pipe = Flux2KleinPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
        # 24GB card vs ~23GB model: offload to fit with headroom for activations.
        pipe.enable_model_cpu_offload()
        self.pipe = pipe
        vol.commit()  # persist the freshly downloaded FLUX weights (no-op if already cached)

    def _gen(self, prompt, image, seed):
        gen = self._torch.Generator("cuda").manual_seed(int(seed))
        result = self.pipe(
            prompt=prompt,
            image=image,
            generator=gen,
            num_inference_steps=STEPS,
            guidance_scale=GUIDANCE,
            height=MAX_SIDE,
            width=MAX_SIDE,
        )
        return result.images[0]

    @modal.method()
    def initial(self, user_photo, art_style: str, scene_prompt: str, seed: int):
        return self._gen(build_initial_prompt(art_style, scene_prompt), [user_photo], seed)

    @modal.method()
    def edit(self, base_image, edit_instruction: str, art_style: str, seed: int):
        return self._gen(build_edit_prompt(art_style, edit_instruction), base_image, seed)