"""Shared LLM tool-use machinery for the demo and chat modes.

A tool-calling LLM (any OpenAI-compatible endpoint) drives the capsule by
translating intent into AICC tool calls. The controller keeps agent-side
working memory (protocol §9): a CURRENT STATE note derived from what the model
has already observed, vision digest for text-only models, and a transcript of
tool calls/results fed back as `tool` messages.
"""

from __future__ import annotations

import base64
import io
import json
import math
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

from aicc.client import AICCClient
from aicc.errors import ToolError
from aicc.protocol import SessionInit, ToolSpec

MISSION = (
    "You are a capsule AI inside a 16x16 m room. Your mission: reach the glowing "
    "beacon pillar and activate it with the interact tool. "
    "World: yaw 0 faces +z, 90 faces +x, coordinates in meters; three wooden "
    "crates block parts of the room. "
    "Perception: proprioception gives position/heading; world_query gives the room "
    "layout and the beacon's exact coordinates (call it once at the start); "
    "vision gives the first-person frame; hear tells you how close you are to the beacon. "
    "Every tool result is prefixed with a CURRENT STATE note: read it — it tells you "
    "your position, heading, distance to the beacon, and the bearing to turn toward it. "
    "Standard loop, one action per turn: "
    "1. call proprioception (and hear every few steps); "
    "2. compute the bearing to the beacon and turn toward it (turn yaw_deg = that bearing, max 90 per call); "
    "3. move forward 1.0 m; if the result says collision, turn 60-90 degrees and go around, "
    "then re-aim. "
    "Never walk into a crate or wall twice in a row. "
    "Call interact ONLY when proprioception puts you within 1.6 m of the beacon. "
    "Keep tool calls short and decisive. Stop once the beacon is active."
)

CHAT_MISSION = (
    "You are the brain of a capsule robot inside a 16x16 m room with a glowing "
    "beacon at the far corner and a few wooden crates. The user talks to you in "
    "natural language (any language) and you translate their intent into tool "
    "calls. "
    "World: yaw 0 faces +z, 90 faces +x, coordinates in meters. "
    "Every tool result is prefixed with a CURRENT STATE note: position, heading, "
    "distance to the beacon and the bearing to turn toward it. "
    "Rules: "
    "1. Take the user's request literally: 'go to the beacon' means turn toward "
    "it and move step by step until close; 'turn left/right' means turn ~90 deg; "
    "'look around' means vision and describe what you see; 'activate' means "
    "interact (only within 1.6 m). "
    "2. One tool call per turn is fine — you will get another turn to continue. "
    "3. If a move reports a collision, turn and go around. "
    "4. When the user asks something not involving movement, just answer briefly. "
    "5. Keep the user informed with one short line after acting (or if you need "
    "a clarification). "
    "Call interact only when within 1.6 m of the beacon."
)

NUDGE = (
    "You have not called a tool. Continue the mission: check the CURRENT STATE note, "
    "then call a tool (turn or move, or interact if within 1.6 m)."
)

COLOR_NAMES = [
    ("beacon_cyan", (120, 210, 235)),
    ("beacon_yellow", (255, 190, 90)),
    ("crate_red", (178, 64, 54)),
    ("crate_blue", (64, 96, 178)),
    ("crate_olive", (128, 128, 60)),
    ("floor_gray", (60, 60, 66)),
    ("wall_gray", (96, 96, 106)),
    ("dark", (30, 30, 36)),
]

LogFn = Callable[[str, str], None]


def _nearest_color(r: float, g: float, b: float) -> str:
    best, best_d = "dark", 1e9
    for name, (cr, cg, cb) in COLOR_NAMES:
        d = (r - cr) ** 2 + (g - cg) ** 2 + (b - cb) ** 2
        if d < best_d:
            best, best_d = name, d
    return best


def vision_digest(png_b64: str, cols: int = 12, rows: int = 9) -> str:
    """Coarse color-grid summary of a vision frame (agent-side preprocessing)."""
    try:
        from PIL import Image

        img = Image.open(io.BytesIO(base64.b64decode(png_b64))).convert("RGB")
    except Exception:  # noqa: BLE001 - digest is best-effort
        return "(no digest: unreadable frame)"
    small = img.resize((cols, rows))
    grid: list[list[tuple[int, int, int]]] = []
    for r in range(rows):
        row: list[tuple[int, int, int]] = []
        for c in range(cols):
            px = small.getpixel((c, r))
            if isinstance(px, (tuple, list)) and len(px) >= 3:
                row.append((int(px[0]), int(px[1]), int(px[2])))
            elif isinstance(px, (int, float)):
                v = int(px)
                row.append((v, v, v))
            else:
                row.append((0, 0, 0))
        grid.append(row)
    lines = [
        " ".join(_nearest_color(*grid[r][c]) for c in range(cols)) for r in range(rows)
    ]
    return f"frame digest ({cols}x{rows}, center = where you face):\n" + "\n".join(
        lines
    )


def to_openai_tools(tools: list[ToolSpec]) -> list[dict[str, Any]]:
    """Convert manifest ToolSpecs to OpenAI function schemas (skip conformance tools)."""
    out = []
    for t in tools:
        if t.id in ("echo", "boom", "bump"):
            continue
        out.append(
            {
                "type": "function",
                "function": {
                    "name": t.id,
                    "description": t.description,
                    "parameters": t.input_schema,
                },
            }
        )
    return out


def ollama_models(base_url: str, api_key: str) -> list[str]:
    """List model ids on an OpenAI-compatible endpoint (best effort)."""
    from openai import AsyncOpenAI

    ids: list[str] = []

    async def _fetch() -> None:
        ac = AsyncOpenAI(base_url=base_url, api_key=api_key)
        try:
            models = await ac.models.list()
        except Exception:  # noqa: BLE001
            return
        ids.extend(m.id for m in models.data if m.id)

    import asyncio

    asyncio.run(_fetch())
    return ids


# ---------------------------------------------------------------------------
# Controller: one conversation with the tool-calling LLM
# ---------------------------------------------------------------------------


@dataclass
class ToolCallResult:
    name: str
    args: dict[str, Any]
    ok: bool
    output: dict[str, Any] | None = None
    error: str | None = None
    digest: str | None = None


@dataclass
class TurnResult:
    """One model round-trip: optional text plus the executed tool calls."""

    text: str
    calls: list[ToolCallResult] = field(default_factory=list)
    interacted: bool = False
    message: str | None = None


def detect_multimodal(base_url: str, model: str) -> bool:
    """Best-effort check whether the model can actually see images.

    For a local ollama we ask /api/show (capabilities include 'vision').
    For other OpenAI-compatible endpoints we cannot introspect — the caller
    may force it with --vision.
    """
    host = base_url.replace("http://", "").replace("https://", "").split("/")[0]
    if host not in ("localhost:11434", "127.0.0.1:11434"):
        return False
    import urllib.request

    endpoint = base_url.rsplit("/v1", 1)[0] + "/api/show"
    try:
        req = urllib.request.Request(
            endpoint,
            data=json.dumps({"model": model}).encode(),
            headers={"Content-Type": "application/json"},
        )
        with urllib.request.urlopen(req, timeout=10) as resp:
            caps = json.loads(resp.read()).get("capabilities", [])
        return "vision" in caps
    except Exception:  # noqa: BLE001 - detection is best-effort
        return False


class LLMController:
    """Drives a tool-calling LLM over an OpenAI-compatible endpoint.

    Holds the message history and the agent-side state (position, heading,
    beacon, path) distilled from tool results. ``invoke()`` runs one model
    turn and executes whatever tools it asks for.
    """

    def __init__(
        self,
        client: AICCClient,
        manifest: SessionInit,
        *,
        base_url: str,
        api_key: str,
        model: str,
        system_prompt: str = MISSION,
        log: LogFn | None = None,
        multimodal: bool | None = None,
    ):
        from openai import AsyncOpenAI

        self.client = client
        self.model = model
        if multimodal is None:
            multimodal = detect_multimodal(base_url, model)
        self.multimodal = multimodal
        self.ac = AsyncOpenAI(base_url=base_url, api_key=api_key)
        self.tools = to_openai_tools(manifest.tools)
        self.messages: list[dict[str, Any]] = [
            {"role": "system", "content": system_prompt},
            {
                "role": "user",
                "content": "Begin. Use the tools to control the capsule.",
            },
        ]
        self.log = log or (lambda role, msg: None)
        self.beacon_pos: tuple[float, float] | None = None
        self.pos: tuple[float, float] | None = None
        self.yaw: float = 0.0
        self.path: list[tuple[float, float]] = []

    # -- agent-side working memory -----------------------------------------

    def state_hint(self) -> str:
        if self.pos is None:
            return ""
        if self.beacon_pos is not None:
            dist = math.hypot(
                self.beacon_pos[0] - self.pos[0], self.beacon_pos[1] - self.pos[1]
            )
            bearing = (
                math.degrees(
                    math.atan2(
                        self.beacon_pos[0] - self.pos[0],
                        self.beacon_pos[1] - self.pos[1],
                    )
                )
                - self.yaw
            ) % 360.0
            return (
                f"CURRENT STATE: you are at ({self.pos[0]:.2f}, {self.pos[1]:.2f}), "
                f"heading {self.yaw:.1f} deg; the beacon is {dist:.2f} m away at relative "
                f"bearing {bearing:.0f} deg. Action: turn yaw_deg={bearing:.0f} to face it, "
                f"then move forward."
            )
        return f"CURRENT STATE: you are at ({self.pos[0]:.2f}, {self.pos[1]:.2f}), heading {self.yaw:.1f} deg."

    def observe(self, output: dict[str, Any]) -> None:
        if isinstance(output.get("position"), dict):
            self.pos = (output["position"]["x"], output["position"]["z"])
        if isinstance(output.get("rotation"), dict):
            self.yaw = output["rotation"]["yaw_deg"] % 360.0
        if self.pos is not None and (not self.path or self.path[-1] != self.pos):
            self.path.append(self.pos)
            if len(self.path) > 4000:
                self.path = self.path[-2000:]

    # -- conversation -------------------------------------------------------

    async def send_user(self, text: str) -> None:
        self.messages.append({"role": "user", "content": text})

    async def invoke(self) -> TurnResult:
        """One model round-trip: get the response and execute its tool calls."""
        try:
            resp = await self.ac.chat.completions.create(
                model=self.model,
                messages=self.messages,
                tools=self.tools,
                tool_choice="auto",
            )
        except Exception as exc:
            self.log("agent", f"LLM error: {exc}")
            raise RuntimeError(f"LLM error: {exc}") from exc

        choice = resp.choices[0]
        text = choice.message.content or ""
        result = TurnResult(text=text)
        calls = choice.message.tool_calls
        if not calls:
            self.messages.append({"role": "assistant", "content": text})
            return result

        self.messages.append(
            {
                "role": "assistant",
                "content": text,
                "tool_calls": [
                    {
                        "id": c.id,
                        "type": "function",
                        "function": {
                            "name": c.function.name,
                            "arguments": c.function.arguments,
                        },
                    }
                    for c in calls
                ],
            }
        )
        for tc in calls:
            name = tc.function.name
            try:
                args = json.loads(tc.function.arguments or "{}")
            except json.JSONDecodeError:
                args = {}
            self.log("tool", f"{name}({json.dumps(args)})")
            try:
                outcome = await self.client.call_tool(name, args)
                output = dict(outcome.output)
                digest = None
                extra = ""
                if name == "vision" and isinstance(output.get("png_b64"), str):
                    digest = vision_digest(output["png_b64"])
                    extra = "\n" + digest
                    output = dict(output, png_b64="<binary, decoded for digest>")
                self.log("bridge", f"ok {json.dumps(output)[:500]}{extra}")
                if name == "world_query":
                    b = output.get("beacon")
                    if isinstance(b, dict):
                        self.beacon_pos = (b["x"], b["z"])
                self.observe(output)
                hint = self.state_hint()
                frame_b64 = None
                if name == "vision" and isinstance(outcome.output.get("png_b64"), str):
                    frame_b64 = outcome.output["png_b64"]
                if self.multimodal and frame_b64 is not None:
                    # Real vision: hand the model the actual frame as an image
                    # (data URI), not as base64 text. Keep the digest too —
                    # it is a cheap textual anchor for the model.
                    content = (hint + "\n" if hint else "") + (
                        f"Camera frame ({output['width']}x{output['height']}, "
                        f"tick {output.get('tick')}) attached as an image.{extra}"
                    )
                    self.messages.append(
                        {"role": "tool", "tool_call_id": tc.id, "content": content}
                    )
                    self.messages.append(
                        {
                            "role": "user",
                            "content": [
                                {
                                    "type": "text",
                                    "text": "This is what your camera sees right now. Use it to orient yourself.",
                                },
                                {
                                    "type": "image_url",
                                    "image_url": {
                                        "url": f"data:image/png;base64,{frame_b64}"
                                    },
                                },
                            ],
                        }
                    )
                else:
                    content = json.dumps(outcome.output) + extra
                    if hint:
                        content = hint + "\n" + content
                    self.messages.append(
                        {"role": "tool", "tool_call_id": tc.id, "content": content}
                    )
                call_result = ToolCallResult(
                    name=name, args=args, ok=True, output=outcome.output, digest=digest
                )
                result.calls.append(call_result)
                if name == "interact" and output.get("success"):
                    result.interacted = True
                    result.message = output.get("message")
            except ToolError as e:
                self.log("bridge", f"error {e.code}: {e.message}")
                self.messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tc.id,
                        "content": json.dumps({"error": e.code, "message": e.message}),
                    }
                )
                result.calls.append(
                    ToolCallResult(name=name, args=args, ok=False, error=e.message)
                )
        return result


# ---------------------------------------------------------------------------
# Autonomous loop (used by the demo; nudges + corrections for weak models)
# ---------------------------------------------------------------------------


async def run_llm_agent_loop(
    controller: LLMController,
    max_steps: int,
    *,
    log: LogFn,
    recorder: Any | None = None,
    nudge_limit: int = 1,
) -> dict[str, Any]:
    """Drive the controller until the mission is done or steps run out."""
    summary: dict[str, Any] = {
        "steps": 0,
        "tool_calls": 0,
        "result": None,
        "interacted": False,
    }
    dist_history: list[float] = []
    corrections = 0
    last_correct_step = -99
    last_collision_step = -99

    def maybe_correct(step: int) -> None:
        nonlocal corrections, last_correct_step
        if controller.pos is None or controller.beacon_pos is None:
            return
        dist = math.hypot(
            controller.beacon_pos[0] - controller.pos[0],
            controller.beacon_pos[1] - controller.pos[1],
        )
        if not math.isfinite(dist) or dist < 2.0:
            return
        dist_history.append(dist)
        recent = dist_history[-4:]
        moving_away = len(recent) == 4 and all(
            recent[i] > recent[i + 1] for i in range(3)
        )
        bearing = (
            math.degrees(
                math.atan2(
                    controller.beacon_pos[0] - controller.pos[0],
                    controller.beacon_pos[1] - controller.pos[1],
                )
            )
            - controller.yaw
        ) % 360.0
        deviation = min(bearing, 360.0 - bearing)
        facing_wrong = dist > 4.0 and deviation > 40.0 and step - last_correct_step > 3
        detouring = step - last_collision_step < 3
        if (moving_away or (facing_wrong and not detouring)) and corrections < 5:
            corrections += 1
            last_correct_step = step
            hint = controller.state_hint()
            msg = (
                "You are not making progress toward the beacon. Stop and follow the CURRENT "
                f"STATE note exactly: {hint} Call turn with that yaw_deg, then move forward."
            )
            controller.messages.append({"role": "user", "content": msg})
            log("agent", "(correction: re-aim at the beacon)")
            dist_history.clear()

    for step in range(max_steps):
        summary["steps"] = step + 1
        if recorder is not None and controller.pos is not None:
            await recorder.snap(controller.client, controller.pos, controller.yaw)
        turn = await controller.invoke()
        if turn.text:
            log("agent", turn.text[:400])
        summary["tool_calls"] += len(turn.calls)
        for c in turn.calls:
            if (
                c.name == "move"
                and c.ok
                and isinstance(c.output, dict)
                and c.output.get("collision")
            ):
                last_collision_step = step
                obj = c.output.get("collision_normal")
                hint_msg = (
                    f"You collided with an obstacle (normal {obj}). Turn yaw_deg=90 and move "
                    "forward twice to get around it, then follow the CURRENT STATE note to re-aim."
                )
                if not (
                    controller.messages
                    and controller.messages[-1].get("content") == hint_msg
                ):
                    controller.messages.append({"role": "user", "content": hint_msg})
                    log("agent", "(collision: go around)")
            maybe_correct(step)
        if turn.interacted:
            summary["interacted"] = True
            summary["result"] = turn.message
            return summary
        if not turn.calls:
            # Nudge a model that drifted into plain text back to acting.
            silent = (not turn.text) or len(turn.text) < 8
            consecutive_nudges = 0
            for msg in reversed(controller.messages):
                if msg.get("content") == NUDGE:
                    consecutive_nudges += 1
                else:
                    break
            if silent or consecutive_nudges >= nudge_limit:
                summary["result"] = "model produced no tool call"
                return summary
            controller.messages.append({"role": "user", "content": NUDGE})
            log("agent", "(nudge: no tool call; continue the mission)")
    summary["result"] = f"exceeded {max_steps} steps"
    return summary