"""MyVillage Activity Instruction API — one FastAPI process, GGUF via llama-cpp-python."""

import json
import logging
import os
import re
import secrets
import threading
import time
import uuid
from contextlib import asynccontextmanager
from typing import Literal

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator

logging.basicConfig(
    level=os.getenv("LOG_LEVEL", "INFO").upper(),
    format="%(asctime)s %(levelname)s %(name)s %(message)s",
)
logger = logging.getLogger("activity-instruction-api")

# Exact evaluated merged checkpoint and its GGUF repository.
SOURCE_MODEL = "mjpsm/activity-instruction-model-500-qwen3.5-2b"
SOURCE_REVISION = "9ee80b07eb127e22c2008ea95fa75c43230b5c0e"
MODEL_REPO = os.getenv("MODEL_REPO", SOURCE_MODEL + "-gguf")
MODEL_FILENAME = os.getenv(
    "MODEL_FILENAME", "activity-instruction-model-500-qwen3.5-2b-Q4_K_M.gguf"
)
MODEL_REVISION = os.getenv("MODEL_REVISION") or None
N_CTX = int(os.getenv("N_CTX", "4096"))
N_THREADS = int(os.getenv("N_THREADS", "2"))
MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "192"))
REPEAT_PENALTY = float(os.getenv("REPEAT_PENALTY", "1.0"))
SEED = int(os.getenv("SEED", "42"))
QUEUE_TIMEOUT = float(os.getenv("QUEUE_TIMEOUT", "90"))

llm: Llama | None = None
model_lock = threading.Lock()

# Exact baseline system prompt from the final 40-case evaluation and GGUF regression.
SYSTEM_INSTRUCTION = '''You are the MyVillage Activity Instruction Model.

You receive the Activity Progression Model output and
the Activity Composition Model output.

Generate clear, actionable instructions for the planned activity.

Rules:
- Follow the progression activity_plan.action and target.
- Address the next_gap.
- Respect demonstrated_state and avoid_repeating.
- Do not assign already-completed work again.
- Do not invent unsupported requirements or details.
- Preserve uncertainty when results are not yet known.
- Return between 1 and 5 concise instructions.
- Return valid JSON only, without explanations.

Required output:
{"output":{"instructions":["instruction 1","instruction 2"]}}'''


class StrictModel(BaseModel):
    model_config = ConfigDict(extra="forbid")


class ActivityPlan(StrictModel):
    action: str = Field(min_length=1)
    target: str = Field(min_length=1)
    avoid_repeating: list[str]


class Progression(StrictModel):
    knowledge_state: Literal["COMPLETED", "PARTIAL", "UNCLEAR", "BLOCKED", "ADVANCED"]
    demonstrated_state: list[str]
    next_gap: str = Field(min_length=1)
    activity_plan: ActivityPlan


class Composition(StrictModel):
    activity_title: str = Field(min_length=1)
    activity_description: str = Field(min_length=1)


class InstructionRequest(StrictModel):
    progression: Progression
    composition: Composition


class InstructionOutput(StrictModel):
    instructions: list[str] = Field(min_length=1, max_length=5)

    @field_validator("instructions")
    @classmethod
    def ensure_nonempty(cls, items: list[str]) -> list[str]:
        if any(not isinstance(v, str) or not v.strip() for v in items):
            raise ValueError("Every instruction must be a nonempty string")
        return items


class InstructionResponse(StrictModel):
    output: InstructionOutput


def build_prompt(payload: dict) -> str:
    """Qwen3.5 two-message chat framing, with its non-thinking assistant prefix.

    Composition uses a raw INPUT/OUTPUT prompt. The Instruction model was
    evaluated as system + user chat, with enable_thinking=False. The original
    Qwen3.5 chat template in that mode opens the assistant with an empty
    <think></think> block. We reproduce that explicitly for create_completion.
    Re-check equivalence against the saved GGUF regressions when deployed.
    """
    user_json = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
    return (
        "<|im_start|>system\n" + SYSTEM_INSTRUCTION + "<|im_end|>\n"
        "<|im_start|>user\n" + user_json + "<|im_end|>\n"
        "<|im_start|>assistant\n<think>\n\n</think>\n\n"
    )


def extract_json(text: str) -> dict:
    cleaned = re.sub(r"^\s*<think>.*?</think>\s*", "", text, flags=re.DOTALL).strip()
    fenced = re.fullmatch(r"```(?:json)?\s*(.*?)\s*```", cleaned, re.DOTALL | re.IGNORECASE)
    if fenced:
        cleaned = fenced.group(1).strip()
    value = json.loads(cleaned)  # Reject arbitrary preamble/trailing non-JSON text.
    if not isinstance(value, dict):
        raise ValueError("Expected a JSON object")
    return value


def load_model() -> Llama:
    token = os.getenv("HF_TOKEN") or None
    logger.info("Downloading GGUF: repo=%s filename=%s revision=%s", MODEL_REPO, MODEL_FILENAME, MODEL_REVISION)
    path = hf_hub_download(
        repo_id=MODEL_REPO,
        filename=MODEL_FILENAME,
        revision=MODEL_REVISION,
        token=token,
    )
    logger.info("Loading model in-process with llama-cpp-python (n_ctx=%s n_threads=%s)", N_CTX, N_THREADS)
    instance = Llama(
        model_path=path,
        n_ctx=N_CTX,
        n_threads=N_THREADS,
        n_gpu_layers=0,  # CPU wheel; same architecture as Composition API.
        seed=SEED,
        verbose=False,
    )
    logger.info("GGUF model loaded successfully")
    return instance


@asynccontextmanager
async def lifespan(app: FastAPI):
    global llm
    llm = load_model()
    try:
        yield
    finally:
        llm = None


app = FastAPI(
    title="MyVillage Activity Instruction API",
    version="1.1.0",
    description="Generate instructions from Progression + Composition using a GGUF via llama-cpp-python.",
    lifespan=lifespan,
)

origins = os.getenv("ALLOWED_ORIGINS", "*").strip()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"] if origins == "*" else [x.strip() for x in origins.split(",") if x.strip()],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.middleware("http")
async def request_logging(request: Request, call_next):
    request_id = request.headers.get("x-request-id") or str(uuid.uuid4())
    start = time.perf_counter()
    if request.url.path in ("/generate", "/predict"):
        required_key = os.getenv("API_KEY")
        if required_key and not secrets.compare_digest(request.headers.get("x-api-key", ""), required_key):
            from fastapi.responses import JSONResponse
            response = JSONResponse(status_code=401, content={"detail": "Invalid or missing API key."})
            response.headers["x-request-id"] = request_id
            return response
    response = await call_next(request)
    response.headers["x-request-id"] = request_id
    logger.info("%s %s status=%s request_id=%s elapsed=%.3fs",
                request.method, request.url.path, response.status_code, request_id, time.perf_counter()-start)
    return response


def generate_instructions(payload: dict) -> InstructionResponse:
    if llm is None:
        raise HTTPException(status_code=503, detail="Model has not loaded yet.")

    prompt = build_prompt(payload)
    acquired = model_lock.acquire(timeout=QUEUE_TIMEOUT)
    if not acquired:
        raise HTTPException(status_code=503, detail="Model is busy. Retry shortly.")
    try:
        result = llm.create_completion(
            prompt=prompt,
            max_tokens=MAX_NEW_TOKENS,
            temperature=0.0,
            repeat_penalty=REPEAT_PENALTY,
            seed=SEED,
            echo=False,
            stop=["<|im_end|>", "<|im_start|>"],
        )
    except Exception as exc:
        logger.exception("llama-cpp-python inference failed")
        raise HTTPException(status_code=500, detail="Model inference failed.") from exc
    finally:
        model_lock.release()

    try:
        choice = result["choices"][0]
        if choice.get("finish_reason") == "length":
            raise ValueError("Generation truncated at max_tokens")
        obj = extract_json(choice["text"])
        return InstructionResponse.model_validate(obj)
    except (KeyError, IndexError, TypeError, ValueError, ValidationError) as exc:
        logger.warning("Model returned invalid instruction JSON/schema: %s", type(exc).__name__)
        raise HTTPException(status_code=502, detail="Model returned invalid or nonconforming instruction JSON.") from exc


@app.get("/")
def root():
    return {
        "name": "MyVillage Activity Instruction API",
        "status": "online" if llm is not None else "starting",
        "model_repo": MODEL_REPO,
        "model_file": MODEL_FILENAME,
        "quantization": "Q4_K_M" if "Q4_K_M" in MODEL_FILENAME.upper() else ("Q8_0" if "Q8_0" in MODEL_FILENAME.upper() else "other"),
        "runtime": "llama-cpp-python",
        "docs": "/docs", "health": "/health", "generate": "/generate",
    }


@app.get("/health")
def health():
    if llm is None:
        raise HTTPException(status_code=503, detail="Model is not loaded yet.")
    return {"status": "ok", "model_loaded": True, "model_repo": MODEL_REPO, "model_file": MODEL_FILENAME}


@app.get("/model")
def model_info():
    return {
        "source_model": SOURCE_MODEL, "source_revision": SOURCE_REVISION,
        "gguf_repo": MODEL_REPO, "gguf_file": MODEL_FILENAME,
        "gguf_revision": MODEL_REVISION,
        "runtime": "llama-cpp-python", "n_ctx": N_CTX, "n_threads": N_THREADS,
        "max_new_tokens": MAX_NEW_TOKENS, "temperature": 0.0,
        "repeat_penalty": REPEAT_PENALTY, "seed": SEED, "thinking": False,
    }


@app.post("/generate", response_model=InstructionResponse)
def generate(request: InstructionRequest):
    return generate_instructions(request.model_dump(mode="json"))


@app.post("/predict", response_model=InstructionResponse, include_in_schema=False)
def predict(request: InstructionRequest):
    return generate(request)