import os
import time

import modal

from sovereign_bench.engine import stream_trial_jsonl
from sovereign_bench.llm import (
    ModelCall,
    ModelResult,
    build_role_messages,
    messages_hash,
)
from sovereign_bench.models import TrialRequest

app = modal.App("sovereign-bench")
GPU_NAME = "H100"
GPU_TIMEOUT_SECONDS = 20 * 60
HF_CACHE_DIR = "/root/.cache/huggingface"

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install("fastapi", "huggingface_hub", "httpx", "pydantic")
    .add_local_dir("sovereign_bench", remote_path="/root/sovereign_bench")
)

model_cache = modal.Volume.from_name("sovereign-bench-model-cache", create_if_missing=True)

vllm_image = (
    modal.Image.from_registry("nvidia/cuda:12.8.1-devel-ubuntu22.04", add_python="3.12")
    .entrypoint([])
    .uv_pip_install(
        "vllm==0.18.1",
        "huggingface_hub[hf_transfer]==0.36.0",
        "transformers",
        "httpx",
        "pydantic",
    )
    .env(
        {
            "HF_HUB_ENABLE_HF_TRANSFER": "1",
            "HF_HOME": HF_CACHE_DIR,
            "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
            "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8": "1",
        }
    )
    .add_local_dir("sovereign_bench", remote_path="/root/sovereign_bench")
)


@app.cls(
    image=vllm_image,
    gpu=GPU_NAME,
    secrets=[modal.Secret.from_name("huggingface")],
    volumes={HF_CACHE_DIR: model_cache},
    timeout=GPU_TIMEOUT_SECONDS,
    scaledown_window=10 * 60,
    max_containers=3,
)
class VllmModel:
    model_id: str = modal.parameter()

    @modal.enter()
    def load(self) -> None:
        from vllm import LLM, SamplingParams

        self.SamplingParams = SamplingParams
        self.llm = LLM(
            model=self.model_id,
            trust_remote_code=True,
            max_model_len=4096,
            gpu_memory_utilization=0.9,
        )

    @modal.method()
    def generate(self, payload: dict) -> dict:
        from sovereign_bench.llm import ModelCallError, clean_model_text

        started = time.perf_counter()
        messages = payload["messages"]
        max_tokens = int(payload.get("max_tokens") or 120)
        temperature = float(payload.get("temperature") or 0.45)
        sampling_params = self.SamplingParams(
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=0.9,
        )
        retry_messages = messages + [
            {
                "role": "user",
                "content": (
                    "Your previous response did not include visible courtroom dialogue. "
                    "Return only the final spoken dialogue now. Do not include <think>, analysis, reasoning, markdown, or notes. /no_think"
                ),
            }
        ]
        last_error: Exception | None = None
        text = ""
        for attempt_messages in (messages, retry_messages):
            outputs = self.llm.chat(
                [attempt_messages],
                sampling_params=sampling_params,
                use_tqdm=False,
                chat_template_kwargs={"enable_thinking": False},
            )
            raw_text = outputs[0].outputs[0].text.strip()
            try:
                text = clean_model_text(raw_text)
                break
            except ModelCallError as exc:
                last_error = exc
        if not text and last_error:
            raise last_error
        return {
            "text": text,
            "latency_ms": int((time.perf_counter() - started) * 1000),
        }


def modal_gpu_enabled() -> bool:
    return os.getenv("SOVEREIGN_DISABLE_MODAL_GPU", "").lower() not in {"1", "true", "yes"}


def modal_gpu_runner(**kwargs) -> ModelResult:
    messages = build_role_messages(
        agent=kwargs["agent"],
        role=kwargs["role"],
        case_summary=kwargs["case_summary"],
        task=kwargs["task"],
        evidence_summary=kwargs["evidence_summary"],
    )
    requested_model = kwargs["model"]
    prompt_hash = messages_hash(messages)

    if modal_gpu_enabled():
        output = VllmModel(model_id=requested_model).generate.remote(
            {
                "messages": messages,
                "max_tokens": kwargs.get("max_tokens", 120),
                "temperature": 0.45,
            }
        )
        return ModelResult(
            text=output["text"],
            input_text="\n\n".join(f"{item.get('role', 'user').upper()}:\n{item.get('content', '')}" for item in messages)
            + "\n\nASSISTANT:\n",
            call=ModelCall(
                model=requested_model,
                provider="modal-gpu-vllm",
                ok=True,
                latency_ms=output["latency_ms"],
                prompt_hash=prompt_hash,
                requested_model=requested_model,
                runtime="modal-gpu-vllm",
                gpu=GPU_NAME,
            ),
        )

    raise RuntimeError("Modal GPU is disabled; no provider fallback is allowed.")


@app.function(image=image, secrets=[modal.Secret.from_name("huggingface")])
def check_huggingface_connection() -> str:
    token = os.getenv("HF_TOKEN")
    if not token:
        return "HF_TOKEN is not available inside Modal."

    from huggingface_hub import HfApi

    user = HfApi(token=token).whoami()["name"]
    return f"Connected to Hugging Face as {user}."


@app.function(
    image=image,
    secrets=[modal.Secret.from_name("huggingface")],
    min_containers=1,
    timeout=GPU_TIMEOUT_SECONDS,
)
@modal.fastapi_endpoint(method="POST", label="trial-stream")
def trial_stream(payload: dict):
    from fastapi.responses import StreamingResponse

    request = TrialRequest.model_validate(payload)
    delay = {"swift": 0.02, "measured": 0.12, "ceremonial": 0.25}[request.speed]
    return StreamingResponse(
        stream_trial_jsonl(request, delay=delay, model_runner=modal_gpu_runner),
        media_type="application/x-ndjson",
    )


@app.local_entrypoint()
def main():
    print(check_huggingface_connection.remote())