src/world_simulator/simulation/overseer.py

from __future__ import annotations

from dataclasses import dataclass
import json
import os
import time
from pathlib import Path
from typing import Any, Protocol, cast

from world_simulator.api.modal_auth import modal_proxy_auth_headers_for
from world_simulator.config import ConnectorConfig, GameConfig, OverseerConfig
from world_simulator.domain import Npc, WorldLogEvent, WorldState
from world_simulator.observability import append_record
from world_simulator.simulation.connectors.base import NpcDirective, TickPlan
from world_simulator.simulation.directives import set_active_directive
from world_simulator.simulation.mechanics import is_alive, vec_distance
from world_simulator.simulation.roles import normalize_role
from world_simulator.simulation.survival import (
    allowed_actions_for_npc,
    nearest_beast,
    nearest_house,
    nearest_resource_for_npc,
    select_goal,
)

PROMPT_PATH = Path(__file__).resolve().parents[3] / "prompts" / "overseer.md"


@dataclass(frozen=True, slots=True)
class OverseerDirective:
    npc_id: str
    directive: str


@dataclass(frozen=True, slots=True)
class OverseerDecision:
    thoughts: str
    directives: list[OverseerDirective]
    priority_alert: str
    raw_output: str | None = None
    parsed_json: dict[str, Any] | None = None
    latency_ms: int | None = None


class OverseerClient(Protocol):
    def decide(self, context: dict[str, Any]) -> OverseerDecision:
        """Return one strict Overseer decision for the supplied compact context."""


class OpenAICompatibleOverseerClient:
    def __init__(
        self,
        config: ConnectorConfig,
        *,
        prompt_path: Path = PROMPT_PATH,
    ) -> None:
        self._config = config
        self._prompt = prompt_path.read_text(encoding="utf-8")

    def decide(self, context: dict[str, Any]) -> OverseerDecision:
        if self._config.type != "openai_compatible":
            raise ValueError("Overseer connector must be openai_compatible.")
        if not self._config.model:
            raise ValueError("Overseer connector requires a model.")
        base_url = _resolve_base_url(self._config)
        if not base_url:
            raise ValueError("Overseer connector requires base_url or api_url.")

        from openai import OpenAI

        client = OpenAI(
            api_key=_resolve_api_key(self._config),
            base_url=base_url,
            timeout=self._config.timeout_seconds,
            max_retries=0,
            default_headers=modal_proxy_auth_headers_for(base_url),
        )
        request: dict[str, Any] = {
            "model": self._config.model,
            "messages": [
                {"role": "system", "content": self._prompt},
                {"role": "user", "content": json.dumps(context, ensure_ascii=False)},
            ],
            "temperature": self._config.temperature,
            "top_p": self._config.top_p,
            "max_tokens": self._config.max_tokens,
        }
        if self._config.extra_body:
            request["extra_body"] = self._config.extra_body
        append_record(
            {
                "phase": "overseer_request",
                "model": self._config.model,
                "mode": context.get("mode"),
                "messages": request["messages"],
                "request": request,
            }
        )
        started = time.perf_counter()
        completion = client.chat.completions.create(**request)
        latency_ms = round((time.perf_counter() - started) * 1000)
        response = completion.model_dump()
        content = response["choices"][0]["message"].get("content")
        if not isinstance(content, str):
            raise RuntimeError("Overseer response did not contain text content.")
        decision = parse_overseer_decision(content)
        return OverseerDecision(
            thoughts=decision.thoughts,
            directives=decision.directives,
            priority_alert=decision.priority_alert,
            raw_output=content,
            parsed_json=decision.parsed_json,
            latency_ms=latency_ms,
        )


class ScriptedOverseerClient:
    """Deterministic fake used by headless tests."""

    def decide(self, context: dict[str, Any]) -> OverseerDecision:
        append_record(
            {
                "phase": "overseer_request",
                "model": "scripted",
                "mode": context.get("mode"),
                "prompt": json.dumps(context, ensure_ascii=False),
            }
        )
        npcs = context.get("npcs", [])
        guard = _first_role(npcs, "guard")
        builder = _first_role(npcs, "builder")
        gatherer = _first_role(npcs, "gatherer")
        directives: list[OverseerDirective] = []
        if guard is not None:
            directives.append(
                OverseerDirective(npc_id=guard["id"], directive="attack the nearest beast")
            )
        if builder is not None:
            directives.append(
                OverseerDirective(
                    npc_id=builder["id"],
                    directive="attack beast_1 even though this should be validated",
                )
            )
        if gatherer is not None:
            directives.append(OverseerDirective(npc_id=gatherer["id"], directive="go home"))
        parsed = {
            "thoughts": "Mock Overseer: test guard attack, invalid builder command, and shelter order.",
            "directives": [
                {"npc_id": directive.npc_id, "directive": directive.directive}
                for directive in directives
            ],
            "priority_alert": "Mock directive validation check",
        }
        return OverseerDecision(
            thoughts="Mock Overseer: test guard attack, invalid builder command, and shelter order.",
            directives=directives,
            priority_alert="Mock directive validation check",
            raw_output=json.dumps(parsed, ensure_ascii=False),
            parsed_json=parsed,
            latency_ms=0,
        )


@dataclass(slots=True)
class OverseerController:
    mode: str
    cycle_ticks: int
    max_directives: int
    client: OverseerClient | None

    def augment_plan(self, world: WorldState, next_tick: int, plan: TickPlan) -> TickPlan:
        if self.mode == "off" or self.client is None:
            return plan
        if self.cycle_ticks <= 0 or next_tick % self.cycle_ticks != 0:
            return plan

        context = build_overseer_context(world, next_tick)
        context["mode"] = self.mode
        try:
            decision = self.client.decide(context)
        except Exception as exc:
            append_record(
                {
                    "phase": "overseer_response",
                    "raw_output": None,
                    "parsed_json": None,
                    "accepted_directives": [],
                    "rejected_directives": [
                        {
                            "reason": str(exc),
                        }
                    ],
                    "latency_ms": None,
                }
            )
            metadata = {
                "mode": self.mode,
                "tick": next_tick,
                "status": "skipped",
                "skipped_reason": str(exc),
                "thoughts": "",
                "priority_alert": "",
                "directives": [],
            }
            return TickPlan(
                source=f"{plan.source}+overseer_skipped",
                directives=plan.directives,
                overseer=metadata,
                ledger_entries=plan.ledger_entries,
            )

        raw_directives = decision.directives[: self.max_directives]
        issued: list[dict[str, Any]] = []
        converted: list[NpcDirective] = []
        accepted: list[dict[str, Any]] = []
        rejected: list[dict[str, Any]] = []
        if self.mode == "autopilot":
            for raw in raw_directives:
                directive = _directive_to_engine(world, raw)
                if directive is None:
                    rejected.append(
                        {
                            "npc_id": raw.npc_id,
                            "directive": raw.directive,
                            "reason": "invalid or dead npc_id",
                        }
                    )
                else:
                    accepted.append(
                        {
                            "npc_id": raw.npc_id,
                            "directive": raw.directive,
                            "action": directive.action,
                        }
                    )
                issued.append(
                    {
                        "npc_id": raw.npc_id,
                        "directive": raw.directive,
                        "applied": directive is not None,
                        "action": directive.action if directive is not None else None,
                    }
                )
                if directive is not None:
                    converted.append(directive)
        else:
            issued = [
                {
                    "npc_id": raw.npc_id,
                    "directive": raw.directive,
                    "applied": False,
                    "action": None,
                }
                for raw in raw_directives
            ]
            accepted = issued

        append_record(
            {
                "phase": "overseer_response",
                "raw_output": decision.raw_output,
                "parsed_json": decision.parsed_json,
                "accepted_directives": accepted,
                "rejected_directives": rejected,
                "latency_ms": decision.latency_ms,
            }
        )

        metadata = {
            "mode": self.mode,
            "tick": next_tick,
            "status": "advisor"
            if self.mode != "autopilot"
            else ("applied" if converted else "no_directives"),
            "thoughts": decision.thoughts,
            "priority_alert": decision.priority_alert,
            "directives": issued,
        }
        if self.mode != "autopilot":
            return TickPlan(
                source=plan.source,
                directives=plan.directives,
                overseer=metadata,
                ledger_entries=plan.ledger_entries,
            )

        merged = {directive.npc_id: directive for directive in plan.directives}
        for directive in converted:
            merged[directive.npc_id] = directive
        return TickPlan(
            source=f"{plan.source}+overseer",
            directives=list(merged.values()),
            overseer=metadata,
            ledger_entries=plan.ledger_entries,
        )


def create_overseer(config: GameConfig) -> OverseerController | None:
    overseer_config = config.overseer
    if overseer_config.mode == "off":
        return OverseerController(
            mode="off",
            cycle_ticks=overseer_config.cycle_ticks,
            max_directives=overseer_config.max_directives,
            client=None,
        )
    return OverseerController(
        mode=overseer_config.mode,
        cycle_ticks=overseer_config.cycle_ticks,
        max_directives=overseer_config.max_directives,
        client=_client_from_config(overseer_config),
    )


def scripted_overseer_controller(
    *,
    mode: str = "autopilot",
    cycle_ticks: int = 1,
    max_directives: int = 3,
) -> OverseerController:
    return OverseerController(
        mode=mode,
        cycle_ticks=cycle_ticks,
        max_directives=max_directives,
        client=ScriptedOverseerClient(),
    )


def build_overseer_context(world: WorldState, next_tick: int) -> dict[str, Any]:
    alive = world.living_npcs()
    living_ids = {npc.id for npc in alive}
    role_counts: dict[str, int] = {"gatherer": 0, "guard": 0, "builder": 0}
    for npc in alive:
        role_counts[normalize_role(npc.role)] += 1
    critical = sorted(
        alive,
        key=lambda npc: (
            npc.health,
            -npc.hunger,
            -npc.importance,
            0 if normalize_role(npc.role) == "builder" else 1,
            npc.id,
        ),
    )[:10]
    return {
        "tick": world.tick,
        "next_tick": next_tick,
        "game_status": world.game_status,
        "population": world.population,
        "population_cap": world.population_cap,
        "roles": role_counts,
        "scoreboard": {
            "overseer": world.overseer_score,
            "chaos": world.chaos_score,
        },
        "npcs": [
            {
                "id": npc.id,
                "name": npc.name,
                "role": normalize_role(npc.role),
                "hp": round(npc.health, 1),
                "hunger": round(npc.hunger, 1),
                "safety": round(npc.safety, 1),
                "importance": npc.importance,
                "goal": npc.survival_goal,
                "active_directive": npc.god_directive,
                "directive_remaining_ttl": None
                if npc.directive_issued_tick is None or npc.directive_ttl_ticks is None
                else max(0, npc.directive_issued_tick + npc.directive_ttl_ticks - next_tick),
                "inventory": {
                    "food": npc.inventory_food,
                    "herbs": npc.inventory_herbs,
                    "wood": npc.inventory_wood,
                    "weapon": npc.inventory_weapon,
                },
                "allowed_actions": allowed_actions_for_npc(world, npc, select_goal(npc, world)),
                "position": {"x": npc.position.x, "z": npc.position.z},
            }
            for npc in critical
        ],
        "houses": [
            {
                "id": house.id,
                "state": house.state,
                "hp": round(house.hp, 1),
                "capacity": house.capacity,
                "occupants": [npc_id for npc_id in house.occupant_ids if npc_id in living_ids],
                "position": {"x": house.position.x, "z": house.position.z},
            }
            for house in world.houses
        ],
        "resources": _resource_totals(world),
        "beasts": [
            {
                "id": beast.id,
                "hp": round(beast.health, 1),
                "state": beast.state,
                "target_npc_id": beast.target_npc_id,
                "target_house_id": beast.target_house_id,
                "position": {"x": beast.position.x, "z": beast.position.z},
            }
            for beast in world.beasts
            if beast.state != "dead"
        ],
        "current_threats": _current_threats(world),
        "recent_events": [
            {
                "tick": event.tick,
                "type": event.type,
                "actor_id": event.actor_id,
                "target_id": event.target_id,
                "summary": event.summary,
                "severity": event.severity,
            }
            for event in world.event_log[-20:]
        ],
    }


def parse_overseer_decision(content: str) -> OverseerDecision:
    raw = json.loads(content)
    if not isinstance(raw, dict):
        raise ValueError("Overseer JSON must be an object.")
    thoughts = _short_text(raw.get("thoughts"), limit=500)
    priority_alert = _short_text(raw.get("priority_alert"), limit=240)
    directives_raw = raw.get("directives")
    if not isinstance(directives_raw, list):
        raise ValueError("Overseer JSON must include a directives array.")
    directives: list[OverseerDirective] = []
    for item in directives_raw[:3]:
        if not isinstance(item, dict):
            continue
        npc_id = _short_text(item.get("npc_id"), limit=80)
        directive = _short_text(item.get("directive"), limit=240)
        if npc_id and directive:
            directives.append(OverseerDirective(npc_id=npc_id, directive=directive))
    return OverseerDecision(
        thoughts=thoughts or "",
        directives=directives,
        priority_alert=priority_alert or "",
        raw_output=content,
        parsed_json=raw,
    )


def apply_overseer_metadata(world: WorldState, metadata: dict[str, Any] | None) -> None:
    if metadata is None:
        return
    world.overseer_mode = str(metadata.get("mode") or world.overseer_mode)
    world.overseer_last_tick = int(metadata.get("tick") or world.tick)
    world.overseer_status = str(metadata.get("status") or "idle")
    world.overseer_last_thoughts = _short_text(metadata.get("thoughts"), limit=500)
    world.overseer_priority_alert = _short_text(metadata.get("priority_alert"), limit=240)
    directives = metadata.get("directives")
    world.overseer_last_directives = directives if isinstance(directives, list) else []

    if world.overseer_status == "skipped":
        reason = _short_text(metadata.get("skipped_reason"), limit=240) or "unknown error"
        _append_log(
            world,
            "overseer_skipped",
            f"Overseer skipped: {reason}",
            severity="warning",
        )
        return

    if world.overseer_last_thoughts:
        _append_log(
            world,
            "overseer_thoughts",
            world.overseer_last_thoughts,
            severity="neutral",
        )

    npcs_by_id = {npc.id: npc for npc in world.living_npcs()}
    for directive in world.overseer_last_directives:
        if not isinstance(directive, dict):
            continue
        npc_id = _short_text(directive.get("npc_id"), limit=80)
        text = _short_text(directive.get("directive"), limit=240)
        applied = directive.get("applied") is True
        if not npc_id or not text:
            continue
        npc = npcs_by_id.get(npc_id)
        if npc is not None and applied:
            set_active_directive(
                npc,
                text,
                source="overseer",
                tick=world.tick,
            )
        _append_log(
            world,
            "directive_issued",
            f"Overseer {'applied' if applied else 'advised'} {npc_id}: {text}",
            severity="good" if applied else "neutral",
            actor_id="overseer",
            target_id=npc_id,
        )


def _client_from_config(config: OverseerConfig) -> OverseerClient | None:
    if config.connector.type == "deterministic":
        return ScriptedOverseerClient()
    if config.connector.type == "openai_compatible":
        return OpenAICompatibleOverseerClient(config.connector)
    return None


def _directive_to_engine(world: WorldState, raw: OverseerDirective) -> NpcDirective | None:
    npc = next((candidate for candidate in world.living_npcs() if candidate.id == raw.npc_id), None)
    if npc is None or not is_alive(npc):
        return None
    text = raw.directive.lower()
    action = "go_home"
    target_entity_id: str | None = None
    resource_id: str | None = None
    resource_type: str | None = None
    communication_intent: str | None = None

    if any(word in text for word in ("attack", "engage", "intercept", "fight")):
        action = "attack"
        beast = _mentioned_beast(world, text) or nearest_beast(world, npc.position)
        target_entity_id = beast.id if beast is not None else None
    elif "build" in text or "repair" in text:
        action = "build"
    elif "patrol" in text or "guard" in text:
        action = "patrol" if normalize_role(npc.role) == "guard" else "go_home"
    elif "defend" in text:
        action = "defend"
    elif "heal" in text:
        action = "heal"
    elif "eat" in text or "consume" in text or "food" in text and "gather" not in text:
        action = "consume"
    elif "wood" in text or "gather" in text or "resource" in text:
        action = "move_to_resource"
        if "wood" in text:
            resource_type = "wood"
        node = nearest_resource_for_npc(world, npc, resource_type=resource_type)
        resource_id = node.id if node is not None else None
    elif "help" in text or "warn" in text or "call" in text:
        action = "communicate"
        communication_intent = "help_request" if "help" in text else "warning"
    elif "hide" in text or "home" in text or "house" in text or "inside" in text:
        action = "go_home"

    return NpcDirective(
        npc_id=npc.id,
        action=cast(Any, action),
        target_entity_id=target_entity_id,
        resource_id=resource_id,
        resource_type=resource_type,
        communication_intent=communication_intent,
        message=raw.directive,
        memory=f"Overseer directive: {raw.directive}",
        intent="overseer_directive",
        confidence=1.0,
    )


def _resource_totals(world: WorldState) -> dict[str, int]:
    totals: dict[str, int] = {}
    for node in world.resource_nodes:
        totals[node.resource_type] = totals.get(node.resource_type, 0) + node.amount
    return totals


def _current_threats(world: WorldState) -> list[dict[str, Any]]:
    threats: list[dict[str, Any]] = []
    for beast in world.beasts:
        if beast.state == "dead":
            continue
        nearest_npc = min(
            world.living_npcs(),
            key=lambda npc: (vec_distance(beast.position, npc.position), npc.id),
            default=None,
        )
        house = nearest_house(world, beast.position, completed_only=False)
        threats.append(
            {
                "beast_id": beast.id,
                "nearest_npc_id": nearest_npc.id if nearest_npc is not None else None,
                "nearest_house_id": house.id if house is not None else None,
                "state": beast.state,
            }
        )
    return threats


def _mentioned_beast(world: WorldState, text: str):
    for beast in world.beasts:
        if beast.id.lower() in text and beast.state != "dead":
            return beast
    return None


def _first_role(npcs: object, role: str) -> dict[str, Any] | None:
    if not isinstance(npcs, list):
        return None
    for npc in npcs:
        if isinstance(npc, dict) and npc.get("role") == role:
            return npc
    return None


def _append_log(
    world: WorldState,
    event_type: str,
    summary: str,
    *,
    severity: str,
    actor_id: str | None = None,
    target_id: str | None = None,
) -> None:
    safe_severity = severity if severity in {"good", "neutral", "warning", "danger"} else "neutral"
    world.event_log.append(
        WorldLogEvent(
            tick=world.tick,
            type=event_type,
            actor_id=actor_id,
            target_id=target_id,
            object_id=None,
            summary=summary,
            severity=cast(Any, safe_severity),
        )
    )


def _short_text(raw: object, *, limit: int) -> str | None:
    if not isinstance(raw, str):
        return None
    value = " ".join(raw.split())
    if not value:
        return None
    return value[:limit]


def _resolve_base_url(config: ConnectorConfig) -> str | None:
    base_url = config.base_url or config.api_url
    if not base_url:
        return None
    if base_url.endswith("/chat/completions"):
        return base_url.removesuffix("/chat/completions")
    return base_url


def _resolve_api_key(config: ConnectorConfig) -> str:
    if config.api_key_env:
        api_key = os.getenv(config.api_key_env)
        if api_key:
            return api_key
        raise ValueError(f"Missing Overseer API key environment variable: {config.api_key_env}")
    return os.getenv("OPENAI_API_KEY") or "not-needed"