"""Record a search mission and render a three-view GIF for the README.

Panes (left -> right):
  top-down debug map (world state: crates, marker, path, capsule)
  first person      (the actual frames the model received)
  chat              (the model's messages, reasoning, and tool calls)

Usage:
    python -m testbed.record_search --out search_mission.gif
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import io
import sys
from pathlib import Path

from aicc.client import AICCClient
from aicc.transport.websocket import WebSocketClientTransport, WebSocketServer
from PIL import Image, ImageDraw, ImageFont

from testbed.bridge import build_bridge
from testbed.llm_agent import (
    SEARCH_MISSION,
    LLMController,
    resolve_provider,
    run_llm_agent_loop,
)

PANEL = 400
ROOM = 16.0
MAX_CHAT_LINES = 21
CHAT_LINE_MAX = 46  # chars per chat line

ROLE_COLORS = {
    "agent": (220, 220, 230),
    "think": (150, 165, 200),
    "tool": (255, 184, 96),
    "bridge": (150, 210, 150),
    "mission": (255, 220, 120),
    "vision": (130, 210, 250),
}


def _font(size: int = 13) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
    try:
        return ImageFont.load_default(size=size)
    except TypeError:  # older Pillow
        return ImageFont.load_default()


def _wrap(text: str, width: int = CHAT_LINE_MAX) -> list[str]:
    words = text.split()
    lines: list[str] = []
    cur = ""
    for w in words:
        if len(cur) + len(w) + 1 > width:
            if cur:
                lines.append(cur)
            cur = w
        else:
            cur = (cur + " " + w).strip()
    if cur:
        lines.append(cur)
    return lines[:3]  # keep the pane readable


class ChatPane:
    """Accumulates the mission chat and renders the last lines."""

    def __init__(self) -> None:
        self.lines: list[tuple[str, str]] = []

    def add(self, role: str, msg: str) -> None:
        for piece in _wrap(msg):
            self.lines.append((role, piece))
        if len(self.lines) > MAX_CHAT_LINES * 2:
            self.lines = self.lines[-MAX_CHAT_LINES:]

    def render(self) -> Image.Image:
        img = Image.new("RGB", (PANEL, PANEL), (12, 12, 18))
        draw = ImageDraw.Draw(img)
        font = _font(13)
        draw.text((8, 6), "chat — what the model says", fill=(160, 170, 190), font=font)
        y = 26
        for role, piece in self.lines[-MAX_CHAT_LINES:]:
            color = ROLE_COLORS.get(role, (200, 200, 200))
            draw.text((8, y), piece[:CHAT_LINE_MAX], fill=color, font=font)
            y += 17
        return img


def _topdown_frame(
    bridge, path: list[tuple[float, float]], step: int, marker_found: bool
) -> Image.Image:
    from testbed.room.render import render_topdown

    img = render_topdown(bridge.world, PANEL)
    draw = ImageDraw.Draw(img)
    scale = PANEL / ROOM
    if len(path) > 1:
        pts = [(x * scale, PANEL - z * scale) for x, z in path]
        draw.line(pts, fill=(255, 170, 60), width=3)
    m = bridge.world.marker_world_pos()
    if m is not None:
        mx, mz = m[0] * scale, PANEL - m[2] * scale
        r = 9
        draw.polygon(
            [(mx, mz - r), (mx - r, mz + r * 0.8), (mx + r, mz + r * 0.8)],
            outline=(255, 255, 255),
            width=2,
        )
    status = f"step {step}  {'FOUND' if marker_found else 'searching...'}"
    draw.text((10, PANEL - 24), status, fill=(255, 220, 120), font=_font(13))
    return img


def _compose(map_img, view_img, chat_img) -> Image.Image:
    canvas = Image.new("RGB", (PANEL * 3, PANEL), (18, 18, 26))
    canvas.paste(map_img, (0, 0))
    canvas.paste(view_img, (PANEL, 0))
    canvas.paste(chat_img, (PANEL * 2, 0))
    draw = ImageDraw.Draw(canvas)
    draw.text((8, 6), "top-down", fill=(160, 170, 190), font=_font(12))
    draw.text((PANEL + 8, 6), "first person", fill=(160, 170, 190), font=_font(12))
    draw.text((PANEL * 2 + 8, 6), "chat", fill=(160, 170, 190), font=_font(12))
    for x in (PANEL, PANEL * 2):
        draw.line([(x, 0), (x, PANEL)], fill=(60, 60, 80))
    return canvas


async def run(args: argparse.Namespace) -> int:
    base_url, api_key, model = resolve_provider(
        provider=args.provider,
        base_url=args.base_url,
        api_key=args.api_key,
        model=args.model,
    )
    bridge = build_bridge()
    # Fixed marker so the GIF tells a clear story (crate_blue, east face).
    bridge.world.marker = {
        "box": "crate_blue",
        "face": "x1",
        "u": 0.5,
        "v": 0.55,
        "size": 0.42,
        "color": (255, 150, 40),
    }
    m = bridge.world.marker_world_pos()
    print(
        f"[record] triangle marker on crate_blue east face at {tuple(round(v, 2) for v in m)}"
    )

    frames: list[Image.Image] = []
    chat = ChatPane()
    marker_seen = False

    def log(role: str, msg: str) -> None:
        print(f"  [{role}] {msg}")
        if role in ("agent", "think", "tool", "mission", "vision"):
            chat.add(role, msg)

    async with WebSocketServer(bridge, port=8765):
        transport = WebSocketClientTransport("ws://127.0.0.1:8765")
        async with AICCClient(transport) as client:
            manifest = await client.handshake()
            ctl = LLMController(
                client,
                manifest,
                base_url=base_url,
                api_key=api_key,
                model=model,
                system_prompt=SEARCH_MISSION,
                log=log,
            )
            ctl.messages = [
                ctl.messages[0],
                {
                    "role": "user",
                    "content": (
                        "MISSION: explore the room, find the orange triangle on the "
                        "back of a crate, and report it. Keep exploring until found."
                    ),
                },
            ]
            chat.add("mission", "MISSION: find the orange triangle")
            orig_frame = ctl._frame_message

            def record_frame(png: str, note: str):
                nonlocal marker_seen
                view = (
                    Image.open(io.BytesIO(base64.b64decode(png)))
                    .convert("RGB")
                    .resize((PANEL, PANEL))
                )
                top = _topdown_frame(bridge, ctl.path, len(frames), marker_seen)
                frames.append(_compose(top, view, chat.render()))
                return orig_frame(png, note)

            ctl._frame_message = record_frame  # type: ignore[method-assign]

            def is_success(turn):
                nonlocal marker_seen
                for c in turn.calls:
                    if (
                        c.name == "report"
                        and c.ok
                        and c.output
                        and c.output.get("verified")
                    ):
                        marker_seen = True
                        return True
                return False

            summary = await run_llm_agent_loop(
                ctl,
                args.steps,
                log=log,
                look_every=1,
                cruise=args.cruise,
                autonomous=True,
                is_success=is_success,
            )

    # Final frame: end state of the map.
    if marker_seen:
        chat.add("mission", f"DONE: {summary.get('result', 'triangle found')}")
    top = _topdown_frame(bridge, ctl.path, len(frames), marker_seen)
    last_view = (
        frames[-1].crop((PANEL, 0, PANEL * 2, PANEL))
        if frames
        else Image.new("RGB", (PANEL, PANEL), (0, 0, 0))
    )
    frames.append(_compose(top, last_view, chat.render()))

    out = Path(args.out)
    if frames:
        frames[0].save(
            out,
            save_all=True,
            append_images=frames[1:],
            duration=args.duration,
            loop=0,
        )
        frames[-1].save(out.with_suffix(".png"))
        print(f"[record] wrote {out} ({len(frames)} frames, {args.duration} ms/frame)")
    print(f"[record] mission result: {summary.get('result')}")
    return 0 if marker_seen else 1


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Record a search mission into a three-view GIF."
    )
    parser.add_argument("--out", default="search_mission.gif")
    parser.add_argument("--provider", default=None)
    parser.add_argument("--base-url", default=None)
    parser.add_argument("--api-key", default=None)
    parser.add_argument("--model", default=None)
    parser.add_argument("--steps", type=int, default=45)
    parser.add_argument("--cruise", type=float, default=0.6)
    parser.add_argument("--duration", type=int, default=350, help="ms per GIF frame")
    args = parser.parse_args()
    try:
        return asyncio.run(run(args))
    except KeyboardInterrupt:
        print("\n[record] interrupted — the partial GIF was NOT written")
        return 130


if __name__ == "__main__":
    sys.exit(main())