"""End-to-end demo: an agent drives the capsule to the beacon and interacts.

Two agents:
  - ``llm`` (default): any tool-calling LLM via an OpenAI-compatible endpoint.
    Ships configured for a local ollama (http://localhost:11434/v1), swap
    ``--base-url``/``--model`` for any provider. The agent perceives through
    the protocol tools only; vision frames are summarized into a coarse color
    grid so text-only models can navigate too. The loop keeps a compact
    CURRENT STATE note (agent-side working memory, protocol §9) and nudges the
    model back on track if it drifts (no tool calls, moving away, collisions).
  - ``scripted``: deterministic sensor-driven fallback (same tools, no LLM).
  - ``auto`` (default): tries the LLM for a bounded number of steps, then
    hands off to the scripted agent so the demo always reaches the beacon.

Every tool call and result is printed to stdout as a transcript. The final
first-person frame is saved to ``--frame`` (default demo_final_frame.png).

Usage:
    python -m testbed.demo --agent llm --model gemma4:e2b
    python -m testbed.demo --agent scripted
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import io
import json
import math
import sys
from pathlib import Path
from typing import Any

from aicc.client import AICCClient
from aicc.protocol import SessionInit
from aicc.transport.websocket import WebSocketClientTransport
from PIL import Image

from testbed.llm_agent import resolve_provider

DEFAULT_URL = "ws://127.0.0.1:8765"


class FrameRecorder:
    """Visual mode: saves first-person frames and a sensor-built top-down map
    with the capsule's path, then stitches a GIF and a summary image.

    Only uses data observed through protocol sensors (world_query,
    proprioception, vision) — the map is reconstructed the way the agent sees
    the world, not read from bridge internals.
    """

    MAP_SIZE = 320
    ROOM = 16.0

    def __init__(self, out_dir: str):
        self.out_dir = Path(out_dir)
        self.out_dir.mkdir(parents=True, exist_ok=True)
        self.layout: dict[str, Any] | None = None
        self.path: list[tuple[float, float]] = []
        self.map_frames: list[Image.Image] = []
        self.view_frames: list[Image.Image] = []
        self.step = 0
        self._log = lambda msg: print(f"[rec] {msg}")

    async def snap(
        self, client: AICCClient, pos: tuple[float, float], yaw: float
    ) -> None:
        """Capture one visual snapshot (one vision call + sensor layout)."""
        if self.layout is None:
            wq = (await client.call_tool("world_query", {})).output
            self.layout = wq
            self._log("world layout captured from sensors")
        self.path.append(pos)
        view = await self._view_frame(client)
        self.view_frames.append(view)
        self.step += 1
        map_img = self._map_frame(pos, yaw)
        self.map_frames.append(map_img)
        view.save(self.out_dir / f"step_{self.step:03d}_view.png")
        map_img.save(self.out_dir / f"step_{self.step:03d}_map.png")

    async def _view_frame(self, client: AICCClient) -> Image.Image:
        from PIL import Image

        out = (await client.call_tool("vision", {})).output
        self._log(f"vision({out['width']}x{out['height']}) frame captured")
        return Image.open(io.BytesIO(base64.b64decode(out["png_b64"]))).convert("RGB")

    def _map_frame(self, pos: tuple[float, float], yaw: float) -> Image.Image:
        from testbed.room.mapview import draw_sensor_map

        return draw_sensor_map(self.layout, pos, yaw, self.path, step=self.step)

    def finish(self, interacted: bool) -> Image.Image:
        """Write demo.gif + demo_summary.png; returns the summary image."""
        from PIL import Image, ImageDraw

        if self.map_frames:
            gif = self.out_dir / "demo.gif"
            self.map_frames[0].save(
                gif,
                save_all=True,
                append_images=self.map_frames[1:],
                duration=350,
                loop=0,
            )
            self._log(f"animation written: {gif}")
        summary = Image.new("RGB", (self.MAP_SIZE + 320, 240), (20, 20, 26))
        if self.map_frames:
            summary.paste(self.map_frames[-1], (0, 0))
        if self.view_frames:
            summary.paste(self.view_frames[-1].resize((320, 240)), (self.MAP_SIZE, 0))
        d = ImageDraw.Draw(summary)
        status = "BEACON ACTIVATED" if interacted else "mission not completed"
        d.text(
            (self.MAP_SIZE + 8, 244 - 16),
            status,
            fill=(255, 220, 120) if interacted else (255, 120, 120),
        )
        out = self.out_dir / "demo_summary.png"
        summary.save(out)
        self._log(f"summary written: {out}")
        return summary


# ---------------------------------------------------------------------------
# LLM agent (OpenAI-compatible)
# ---------------------------------------------------------------------------


async def run_llm_agent(
    client: AICCClient,
    manifest: SessionInit,
    *,
    base_url: str,
    api_key: str,
    model: str,
    max_steps: int,
    recorder: FrameRecorder | None = None,
    vision: bool | None = None,
    digest: bool | None = None,
    look_every: int = 0,
    cruise: float = 0.0,
    autonomous: bool = False,
) -> dict[str, Any]:
    """Autonomous LLM run: controller + nudge/correct loop (see llm_agent)."""
    from testbed.llm_agent import LLMController, run_llm_agent_loop

    controller = LLMController(
        client,
        manifest,
        base_url=base_url,
        api_key=api_key,
        model=model,
        log=lambda role, msg: print(f"[{role}] {msg}"),
        multimodal=vision,
        digest=digest,
    )
    try:
        return await run_llm_agent_loop(
            controller, max_steps, log=print_log, recorder=recorder
        )
    except RuntimeError as exc:
        return {"steps": 0, "tool_calls": 0, "result": str(exc), "interacted": False}


def print_log(role: str, msg: str) -> None:
    print(f"[{role}] {msg}")


# ---------------------------------------------------------------------------
# Scripted agent (sensor-driven fallback, no LLM)
# ---------------------------------------------------------------------------


async def run_scripted_agent(
    client: AICCClient,
    manifest: SessionInit,
    max_steps: int,
    recorder: FrameRecorder | None = None,
) -> dict[str, Any]:
    summary: dict[str, Any] = {
        "steps": 0,
        "tool_calls": 0,
        "result": None,
        "interacted": False,
    }

    def log(role: str, msg: str) -> None:
        print(f"[{role}] {msg}")

    wq = (await client.call_tool("world_query", {})).output
    beacon = wq["beacon"]
    bx, bz = beacon["x"], beacon["z"]
    log("agent", f"mission: reach beacon at ({bx}, {bz})")

    consecutive_collisions = 0
    detour_steps = 0  # >0: escaping an obstacle before re-aiming at the beacon
    detour_turns = 0
    for step in range(max_steps):
        summary["steps"] = step + 1
        prop = (await client.call_tool("proprioception", {})).output
        summary["tool_calls"] += 1
        pos = prop["position"]
        if recorder is not None:
            await recorder.snap(
                client, (pos["x"], pos["z"]), prop["rotation"]["yaw_deg"]
            )
        dist = math.hypot(bx - pos["x"], bz - pos["z"])
        log(
            "agent",
            f"at ({pos['x']:.2f}, {pos['z']:.2f}) heading {prop['rotation']['yaw_deg']:.1f} deg, {dist:.2f} m to beacon",
        )

        if dist <= 1.5:
            la = (await client.call_tool("look_at", {"target": "beacon"})).output
            summary["tool_calls"] += 1
            log("tool", 'look_at({"target": "beacon"})')
            log("bridge", f"ok {json.dumps(la)}")
            frame = (await client.call_tool("vision", {})).output
            summary["tool_calls"] += 1
            log("tool", "vision({})")
            log("bridge", f"ok {frame['width']}x{frame['height']} frame captured")
            res = (await client.call_tool("interact", {})).output
            summary["tool_calls"] += 1
            log("tool", 'interact({"target": "beacon"})')
            log("bridge", f"ok {json.dumps(res)}")
            if res["success"]:
                summary["interacted"] = True
                summary["result"] = res["message"]
                return summary
            log("agent", "not in reach yet; continue")

        if detour_steps > 0:
            # Escape the obstacle before re-aiming: keep the detour heading.
            mv = (await client.call_tool("move", {"forward": 0.8, "duration": 0.6})).output
            summary["tool_calls"] += 1
            log("tool", 'move({"forward": 0.8, "duration": 0.6}) [detour]')
            log("bridge", f"ok {json.dumps(mv)}")
            if mv.get("moved", 0.0) < 0.2 and detour_turns < 6:
                detour_turns += 1
                t = (await client.call_tool("turn", {"yaw_deg": 45.0})).output
                summary["tool_calls"] += 1
                log("tool", 'turn({"yaw_deg": 45.0}) [detour]')
                log("bridge", f"ok {json.dumps(t)}")
            else:
                detour_steps -= 1
            continue

        target_yaw = math.degrees(math.atan2(bx - pos["x"], bz - pos["z"])) % 360.0
        cur_yaw = prop["rotation"]["yaw_deg"] % 360.0
        delta = (target_yaw - cur_yaw + 540.0) % 360.0 - 180.0

        # Keep off the walls so the capsule can round corners.
        if pos["x"] < 1.2:
            delta = min(delta, -60.0)
        elif pos["x"] > 15.0:
            delta = max(delta, 60.0)
        if pos["z"] < 1.2:
            delta = max(delta, 60.0)
        elif pos["z"] > 15.0:
            delta = min(delta, -60.0)

        if abs(delta) > 4.0:
            turn = (await client.call_tool("turn", {"yaw_deg": delta * 0.8})).output
            summary["tool_calls"] += 1
            log("tool", f'turn({{"yaw_deg": {delta * 0.8:.1f}}})')
            log("bridge", f"ok {json.dumps(turn)}")
        else:
            mv = (await client.call_tool("move", {"forward": 0.8, "duration": 0.6})).output
            summary["tool_calls"] += 1
            log("tool", 'move({"forward": 0.8, "duration": 0.6})')
            log("bridge", f"ok {json.dumps(mv)}")
            if mv.get("collision"):
                consecutive_collisions += 1
                angle = 40.0 if consecutive_collisions < 4 else 90.0
                t = (await client.call_tool("turn", {"yaw_deg": angle})).output
                summary["tool_calls"] += 1
                log("tool", f'turn({{"yaw_deg": {angle}}})')
                log("bridge", f"ok {json.dumps(t)}")
                detour_steps = 3
                detour_turns = 0
            else:
                consecutive_collisions = 0
    summary["result"] = f"exceeded {max_steps} steps"
    return summary


# ---------------------------------------------------------------------------
# Shared runner
# ---------------------------------------------------------------------------


async def run_demo(args: argparse.Namespace) -> dict[str, Any]:
    transport = WebSocketClientTransport(args.url)
    async with AICCClient(transport) as client:
        manifest = await client.handshake()
        print(f"[handshake] session {manifest.session_id} world {manifest.world.name}")
        print(f"[handshake] tools: {[t.id for t in manifest.tools]}\n")

        if args.frames_dir:
            recorder = FrameRecorder(args.frames_dir)
        else:
            recorder = None

        if args.agent == "llm":
            summary = await run_llm_agent(
                client,
                manifest,
                base_url=args.base_url,
                api_key=args.api_key,
                model=args.model,
                max_steps=args.max_steps,
                recorder=recorder,
                vision=args.vision,
                digest=args.digest,
                look_every=args.look_every,
                cruise=args.cruise,
                autonomous=args.free,
            )
        elif args.agent == "scripted":
            summary = await run_scripted_agent(
                client, manifest, args.max_steps, recorder=recorder
            )
        else:  # auto
            from openai import AsyncOpenAI

            probe = AsyncOpenAI(base_url=args.base_url, api_key=args.api_key)
            try:
                await asyncio.wait_for(probe.models.list(), timeout=5.0)
                llm_steps = min(args.max_steps, 15)
                summary = await run_llm_agent(
                    client,
                    manifest,
                    base_url=args.base_url,
                    api_key=args.api_key,
                    model=args.model,
                    max_steps=llm_steps,
                    recorder=recorder,
                    vision=args.vision,
                    digest=args.digest,
                    look_every=args.look_every,
                    cruise=args.cruise,
                    autonomous=args.free,
                )
                if summary.get("interacted"):
                    return summary
                print(
                    f"\n[auto] LLM did not finish in {llm_steps} steps "
                    f"({summary.get('result')}); handing off to the scripted agent\n"
                )
                summary = await run_scripted_agent(
                    client, manifest, args.max_steps, recorder=recorder
                )
            except Exception as exc:  # noqa: BLE001 - fall back to scripted
                print(
                    f"\n[auto] LLM endpoint unavailable ({type(exc).__name__}: {exc}); falling back to scripted agent"
                )
                summary = await run_scripted_agent(
                    client, manifest, args.max_steps, recorder=recorder
                )

        if recorder is not None:
            recorder.finish(summary.get("interacted", False))

        frame = (await client.call_tool("vision", {})).output
        if args.frame:
            raw = base64.b64decode(frame["png_b64"])

            def _write_frame() -> None:
                with open(args.frame, "wb") as fh:
                    fh.write(raw)

            await asyncio.to_thread(_write_frame)
            print(f"\n[final frame saved] {args.frame}")
        print(
            f"\n[demo done] steps={summary['steps']} tool_calls={summary['tool_calls']} "
            f"interacted={summary['interacted']} result={summary['result']}"
        )
        return summary


def main() -> int:
    parser = argparse.ArgumentParser(
        description="AICC capsule demo: LLM or scripted agent drives the capsule to the beacon."
    )
    parser.add_argument(
        "--url",
        default=DEFAULT_URL,
        help=f"bridge WebSocket URL (default {DEFAULT_URL})",
    )
    parser.add_argument(
        "--agent",
        choices=["llm", "scripted", "auto"],
        default="auto",
        help="agent driver (default auto: LLM if reachable, else scripted)",
    )
    parser.add_argument(
        "--base-url",
        default=None,
        help="OpenAI-compatible endpoint (default per --provider)",
    )
    parser.add_argument(
        "--provider",
        choices=["polza", "openai", "ollama"],
        default=None,
        help="provider preset: endpoint + key from env (POLZA_API_KEY/OPENAI_API_KEY)",
    )
    parser.add_argument(
        "--model",
        default=None,
        help="model id on the endpoint (default per provider: polza -> openai/gpt-5.6-luna)",
    )
    parser.add_argument(
        "--api-key",
        default=None,
        help="API key (default: $POLZA_API_KEY / $OPENAI_API_KEY per provider)",
    )
    parser.add_argument("--max-steps", type=int, default=60, help="max agent steps")
    parser.add_argument(
        "--vision",
        action=argparse.BooleanOptionalAction,
        default=None,
        help="pass real camera frames to the model as images (auto-detected for local ollama)",
    )
    parser.add_argument(
        "--digest",
        action=argparse.BooleanOptionalAction,
        default=None,
        help="always include the color-grid digest alongside images (off by default for multimodal models)",
    )
    parser.add_argument(
        "--look-every",
        type=int,
        default=3,
        help="attach a fresh camera frame every N agent steps (0 disables; default 3)",
    )
    parser.add_argument(
        "--cruise",
        type=float,
        default=1.0,
        help="proactive motion: glide forward (meters) while the LLM thinks (0 disables)",
    )
    parser.add_argument(
        "--free",
        action="store_true",
        help="fewer restrictions: no re-aim corrections or collision hints; the model plans freely",
    )
    parser.add_argument(
        "--frame",
        default="demo_final_frame.png",
        help="where to save the final vision frame",
    )
    parser.add_argument(
        "--frames-dir",
        default=None,
        help="visual mode: save first-person + top-down map frames each step "
        "into this directory, plus demo.gif animation and demo_summary.png",
    )
    args = parser.parse_args()
    try:
        args.base_url, args.api_key, args.model = resolve_provider(
            provider=args.provider,
            base_url=args.base_url,
            api_key=args.api_key,
            model=args.model,
        )
        if args.provider:
            print(f"[demo] provider: {args.provider} -> {args.base_url}  model: {args.model}")
    except ValueError as exc:
        print(f"[demo] {exc}")
        return 1
    try:
        summary = asyncio.run(run_demo(args))
    except KeyboardInterrupt:
        return 2
    except Exception as exc:  # noqa: BLE001 - friendly failure instead of a traceback
        print(f"\n[demo] failed: {type(exc).__name__}: {exc}")
        print("[demo] is the bridge running?  scripts/run_bridge.sh")
        return 1
    return 0 if summary.get("interacted") else 1


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