server/game_api.py

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from pathlib import Path
from threading import Lock
from typing import Any

from fastapi import APIRouter
from pydantic import BaseModel, Field

from models import BDOAction, BDOObservation
from server.bdo_environment import BDOEnvironment

_logger = logging.getLogger(__name__)

# Path where train_unsloth.py saves the LoRA adapter after GRPO training.
_MODEL_DIR = Path(__file__).resolve().parents[1] / "artifacts" / "grpo_model"

# Module-level cache: populated on first _policy() call if the adapter exists.
_llm_cache: dict[str, Any] | None = None
_llm_load_attempted: bool = False


def _extract_json_object(text: str) -> str:
    """Return the first complete {...} block found in text."""
    start = text.find("{")
    end   = text.rfind("}")
    if start == -1 or end == -1 or end < start:
        raise ValueError("No JSON object found in model output.")
    return text[start : end + 1]


def _try_load_llm() -> dict[str, Any] | None:
    """
    Attempt to load the trained LoRA adapter from artifacts/grpo_model/.
    Returns a dict with 'model', 'tokenizer', 'device' on success, else None.
    Falls back silently so the server always starts even without the adapter.
    """
    global _llm_cache, _llm_load_attempted
    if _llm_load_attempted:
        return _llm_cache
    _llm_load_attempted = True

    if not _MODEL_DIR.exists():
        _logger.info("No trained adapter at %s — using heuristic policy.", _MODEL_DIR)
        return None

    try:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
        from peft import PeftModel

        _logger.info("Loading trained BDO adapter from %s …", _MODEL_DIR)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype  = torch.float16 if device == "cuda" else torch.float32

        tokenizer = AutoTokenizer.from_pretrained(str(_MODEL_DIR), trust_remote_code=True)

        base = AutoModelForCausalLM.from_pretrained(
            "Qwen/Qwen2.5-1.5B-Instruct",
            torch_dtype=dtype,
            device_map="auto" if device == "cuda" else None,
            trust_remote_code=True,
        )
        model = PeftModel.from_pretrained(base, str(_MODEL_DIR))
        if device == "cpu":
            model = model.to(device)
        model.eval()

        _llm_cache = {"model": model, "tokenizer": tokenizer, "device": device}
        _logger.info("Trained adapter loaded on %s. UI will use LLM policy.", device)
        return _llm_cache

    except Exception as exc:
        _logger.warning("Could not load trained adapter (%s) — falling back to heuristic.", exc)
        return None


class ResetRequest(BaseModel):
    scenario: str | None = None
    seed: int | None = None


@dataclass
class UISimulator:
    env: BDOEnvironment = field(default_factory=BDOEnvironment)
    lock: Lock = field(default_factory=Lock)
    initialized: bool = False
    cumulative_reward: float = 0.0
    last_observation: BDOObservation | None = None
    recent_actions: list[dict[str, Any]] = field(default_factory=list)

    def _heuristic_policy(self, observation: dict[str, Any]) -> dict[str, Any]:
        """Rule-based fallback used when no trained adapter is available."""
        nodes = observation["nodes"]
        highest_demand = max(nodes, key=lambda node: node["reported_demand"])
        weakest_signal = min(nodes, key=lambda node: node["biometric_signal"])
        queue = observation.get("high_risk_queue", [])
        shocks = set(observation["meta"].get("active_shocks", []))

        actions: list[dict[str, Any]] = []

        if queue:
            actions.append(
                {"name": "reject_transfer", "params": {"transfer_id": queue[0]["transfer_id"]}}
            )

        if weakest_signal["biometric_signal"] < 0.52:
            actions.append(
                {"name": "dispatch_repair", "params": {"village": weakest_signal["village"]}}
            )
        elif weakest_signal["biometric_signal"] < 0.68:
            actions.append(
                {
                    "name": "trigger_field_audit",
                    "params": {"village": weakest_signal["village"]},
                }
            )

        if shocks:
            target = next((node for node in nodes if node["village"] in shocks), highest_demand)
            actions.append({"name": "issue_early_warning", "params": {"village": target["village"]}})
            actions.append({"name": "deploy_reserve", "params": {"village": target["village"], "amount": 2000}})
        else:
            reserve_build = min(1500, max(500, observation["treasury"]["district_budget"] // 12))
            actions.append({"name": "build_reserve", "params": {"amount": reserve_build}})

        spend = min(
            observation["treasury"]["district_budget"],
            max(2200, int(highest_demand["reported_demand"] * 0.90)),
        )
        actions.append(
            {"name": "allocate_funds", "params": {"village": highest_demand["village"], "amount": spend}}
        )
        mode = "conservative" if queue or weakest_signal["biometric_signal"] < 0.72 else "permissive"
        actions.append({"name": "approve_batch", "params": {"village": highest_demand["village"], "mode": mode}})

        average_signal = sum(node["biometric_signal"] for node in nodes) / max(1, len(nodes))
        predicted_fraud = round(min(0.95, max(0.08, 1.0 - average_signal)), 3)
        return {
            "predicted_fraud_level": predicted_fraud,
            "thought_process": (
                f"Prioritize {highest_demand['village']} demand; repair or audit "
                f"{weakest_signal['village']} sensors; actively suppress queue risk."
            ),
            "actions": actions,
        }

    def _llm_policy(self, observation: dict[str, Any], llm: dict[str, Any]) -> dict[str, Any]:
        """Call the trained Qwen adapter to generate the next action batch."""
        from bdo_ai_env.training import build_prompt

        model     = llm["model"]
        tokenizer = llm["tokenizer"]
        device    = llm["device"]

        prompt  = build_prompt(observation)
        inputs  = tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=3584,
        )
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with model.generation_config.__class__.__init__.__func__.__globals__.get(
            "__builtins__", {}
        ) if False else __import__("contextlib").nullcontext():
            import torch
            with torch.no_grad():
                outputs = model.generate(
                    **inputs,
                    max_new_tokens=512,
                    do_sample=False,
                    use_cache=True,
                    pad_token_id=tokenizer.eos_token_id,
                )

        completion = tokenizer.decode(
            outputs[0][inputs["input_ids"].shape[-1]:],
            skip_special_tokens=True,
        )
        try:
            return BDOAction.model_validate_json(
                _extract_json_object(completion)
            ).model_dump(mode="json", exclude_none=True)
        except Exception:
            # Malformed output — fall back to heuristic for this step
            _logger.warning("LLM output could not be parsed; using heuristic for this step.")
            return self._heuristic_policy(observation)

    def _policy(self, observation: dict[str, Any]) -> dict[str, Any]:
        """Route to LLM policy if trained adapter is available, else heuristic."""
        llm = _try_load_llm()
        if llm is not None:
            return self._llm_policy(observation, llm)
        return self._heuristic_policy(observation)

    @staticmethod
    def _serialize_observation(observation: BDOObservation) -> dict[str, Any]:
        payload = observation.model_dump(mode="json", exclude_none=True)
        info = observation.info or observation.metadata or {}
        payload["info"] = info
        return payload

    def reset(self, scenario: str | None = None, seed: int | None = None) -> dict[str, Any]:
        with self.lock:
            self.env = BDOEnvironment(scenario=scenario, seed=seed)
            self.last_observation = self.env.reset(seed=seed, scenario=scenario)
            self.initialized = True
            self.cumulative_reward = 0.0
            self.recent_actions = []
            return {
                "observation": self._serialize_observation(self.last_observation),
                "done": bool(self.last_observation.done),
                "cumulative_reward": self.cumulative_reward,
                "last_action": None,
                "recent_actions": self.recent_actions,
                "ground_truth": self._get_ground_truth(),
                "agent_mode": "llm" if _llm_cache else "heuristic",
            }

    def state(self) -> dict[str, Any]:
        with self.lock:
            if not self.initialized or self.last_observation is None:
                return self.reset()
            return {
                "observation": self._serialize_observation(self.last_observation),
                "done": bool(self.last_observation.done),
                "cumulative_reward": round(self.cumulative_reward, 4),
                "last_action": self.recent_actions[-1] if self.recent_actions else None,
                "recent_actions": self.recent_actions[-6:],
                "ground_truth": self._get_ground_truth(),
                "agent_mode": "llm" if _llm_cache else "heuristic",
            }

    def _get_ground_truth(self) -> dict[str, Any]:
        """Extract hidden world state for God Mode visualization."""
        try:
            hidden = self.env.world.world
            villages_gt: dict[str, Any] = {}
            for name, v in hidden.villages.items():
                villages_gt[name] = {
                    "hardware_health": round(float(v.hardware_health), 3),
                    "sensor_reliability": round(float(v.sensor_reliability), 3),
                    "actual_fraud_level": round(float(v.actual_fraud_level), 3),
                    "actual_demand": int(v.actual_demand),
                    "wealth": round(float(v.wealth), 3),
                    "stability_index": round(float(v.stability_index), 3),
                    "latent_regime": str(v.latent_regime),
                    "audit_active": bool(v.audit_active),
                    "repair_pending_months": int(v.repair_pending_months),
                    "unmet_demand_ratio": round(float(v.unmet_demand_ratio), 3),
                    "funds_leaked_to_sybils": int(v.funds_leaked_to_sybils),
                }
            return {
                "month": int(hidden.month),
                "shock_month": int(hidden.shock_month),
                "predicted_fraud_level": round(float(hidden.predicted_fraud_level), 3),
                "villages": villages_gt,
            }
        except Exception:
            return {}

    def simulate(self) -> dict[str, Any]:
        with self.lock:
            if not self.initialized or self.last_observation is None:
                self.reset()

            assert self.last_observation is not None
            if self.last_observation.done:
                self.reset()

            assert self.last_observation is not None
            current = self.last_observation.model_dump(mode="json", exclude_none=True)
            action_payload = self._policy(current)
            action = BDOAction.model_validate(action_payload)
            next_observation = self.env.step(action)
            reward = float(next_observation.reward or 0.0)
            self.cumulative_reward += reward
            self.last_observation = next_observation

            info = next_observation.info or next_observation.metadata or {}
            action_record = {
                "month": current["meta"]["month"],
                "thought_process": action_payload.get("thought_process", ""),
                "actions": action_payload.get("actions", []),
                "executed_actions": info.get("executed_actions", []),
                "training_reward": info.get("training_reward", reward),
                "reward": reward,
            }
            self.recent_actions.append(action_record)
            self.recent_actions = self.recent_actions[-24:]

            return {
                "observation": self._serialize_observation(next_observation),
                "done": bool(next_observation.done),
                "cumulative_reward": round(self.cumulative_reward, 4),
                "last_action": action_record,
                "recent_actions": self.recent_actions[-6:],
                "ground_truth": self._get_ground_truth(),
                "agent_mode": "llm" if _llm_cache else "heuristic",
            }


router = APIRouter(prefix="/api/ui", tags=["ui"])
simulator = UISimulator()


@router.get("/state")
def get_state() -> dict[str, Any]:
    return simulator.state()


@router.post("/reset")
def reset(payload: ResetRequest | None = None) -> dict[str, Any]:
    payload = payload or ResetRequest()
    return simulator.reset(scenario=payload.scenario, seed=payload.seed)


@router.post("/simulate")
def simulate() -> dict[str, Any]:
    return simulator.simulate()