"""Interactive chat mode: talk to the model, it drives the capsule.

Type natural-language commands (any language) and the LLM translates them
into AICC tool calls: 'go to the beacon', 'turn left', 'look around',
'activate the beacon'. Every tool call and result is printed as a transcript,
the current vision frame is saved to ``chat_frame.png`` and the sensor-built
top-down map to ``chat_map.png`` after each turn. Works great together with
the live viewer (scripts/run_live.sh) — watch the capsule in your browser
while you talk to it.

Commands:
    /state      print current position/heading (proprioception)
    /look       take a vision frame and save it
    /map        save the top-down map (chat_map.png)
    /models     list models on the endpoint
    /model X    switch the model mid-session
    /steps N    auto-continue budget for 'go to X' style requests
    /help       this text
    /exit       quit

Usage:
    python -m testbed.chat [--model gemma4:e2b] [--model gemma4:12b]
    python -m testbed.chat --base-url https://api.openai.com/v1 --model gpt-4o-mini
"""

from __future__ import annotations

import argparse
import asyncio
import sys

from aicc.client import AICCClient
from aicc.transport.websocket import WebSocketClientTransport

from testbed.llm_agent import (
    CHAT_MISSION,
    LLMController,
    ollama_models,
    resolve_provider,
)

DEFAULT_URL = "ws://127.0.0.1:8765"

HELP = """\
You are talking to the capsule's brain (LLM). Type what you want, e.g.:
  'иди к маяку' / 'go to the beacon'   — the model walks there step by step
  'повернись налево' / 'turn left'     — quick turn
  'осмотрись' / 'look around'          — vision + description
  'активируй маяк' / 'activate it'     — interact (must be close)
Missions (autonomous goal, keeps trying until done):
  /mission             — goal: reach and activate the beacon
  /mission <text>      — any goal, e.g. /mission дойди до маяка
  /status              — mission progress + capsule state
  /stop                — cancel the mission (REPL stays usable)
Other commands: /state /look /map /models /model X /steps N /help /exit"""


async def chat_loop(client: AICCClient, manifest, args: argparse.Namespace) -> int:
    model = args.model
    controller = LLMController(
        client,
        manifest,
        base_url=args.base_url,
        api_key=args.api_key,
        model=model,
        system_prompt=CHAT_MISSION,
        log=lambda role, msg: print(f"  [{role}] {msg}"),
        multimodal=args.vision,
        digest=args.digest,
    )
    auto_steps = args.auto_steps
    print(f"[chat] model: {model}  (endpoint {args.base_url})")
    if controller.multimodal:
        print("[chat] vision: ON — the model sees the actual camera frames")
    else:
        print("[chat] vision: text-only model — frames are sent as a color-grid digest")
    print("[chat] type your commands; /help for the command list; /exit to quit\n")

    async def cmd_state() -> None:
        out = (await client.call_tool("proprioception", {})).output
        print(
            f"  pos ({out['position']['x']:.2f}, {out['position']['z']:.2f})  "
            f"heading {out['rotation']['yaw_deg']:.1f} deg  health {out['health']}"
        )

    async def cmd_look() -> None:
        out = (await client.call_tool("vision", {})).output
        import base64

        raw = base64.b64decode(out["png_b64"])

        def _write() -> None:
            with open("chat_frame.png", "wb") as fh:
                fh.write(raw)

        await asyncio.to_thread(_write)
        print(f"  frame saved to chat_frame.png ({out['width']}x{out['height']})")

    def save_map(quiet: bool = False) -> None:
        if controller.pos is None:
            if not quiet:
                print("  (no position yet — ask the capsule to move first)")
            return
        from testbed.room.mapview import draw_sensor_map

        layout = {"obstacles": [], "beacon": {"x": 12.5, "z": 12.5}}
        if controller.beacon_pos is not None:
            layout["beacon"] = {
                "x": controller.beacon_pos[0],
                "z": controller.beacon_pos[1],
            }
        img = draw_sensor_map(layout, controller.pos, controller.yaw, controller.path)
        img.save("chat_map.png")
        if not quiet:
            print(
                f"  map saved to chat_map.png (pos {controller.pos[0]:.1f}, {controller.pos[1]:.1f})"
            )

    async def cmd_models() -> None:
        print("  fetching models…")
        for m in ollama_models(args.base_url, args.api_key):
            print(f"    {m}")

    # -- mission: autonomous goal run in the background ---------------------
    # The REPL stays responsive: /status to check progress, /stop to cancel.
    mission: asyncio.Task | None = None
    mission_goal = "reach the beacon and activate it"
    user_turns = 0

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

    async def map_saver() -> None:
        while True:
            try:
                save_map(quiet=True)
            except Exception as exc:  # noqa: BLE001 - never let the saver die
                print(f"  [map] save error: {type(exc).__name__}: {exc}")
            await asyncio.sleep(3.0)

    async def mission_runner(goal_text: str) -> None:
        from testbed.llm_agent import run_llm_agent_loop

        if args.search:
            from testbed.llm_agent import SEARCH_MISSION

            controller.messages = [
                {"role": "system", "content": SEARCH_MISSION},
                {
                    "role": "user",
                    "content": (
                        f"MISSION: {goal_text}. "
                        "Keep exploring with tools until the triangle is found and reported."
                    ),
                },
            ]
        else:
            controller.messages.append(
                {
                    "role": "user",
                    "content": (
                        f"MISSION (set by the user): {goal_text}. "
                        "Keep calling tools and do not stop until the goal is achieved. "
                        "Report only when done."
                    ),
                }
            )
        saver = asyncio.create_task(map_saver())
        try:
            for attempt in range(1, args.mission_retries + 2):
                print(
                    f"[mission] attempt {attempt}/{args.mission_retries + 1}: {goal_text}"
                )
                try:
                    summary = await run_llm_agent_loop(
                        controller,
                        args.mission_steps,
                        log=mission_log,
                        nudge_limit=args.mission_nudges,
                        look_every=args.look_every,
                        cruise=args.cruise,
                        autonomous=args.free or args.search,
                        is_success=(
                            (
                                lambda turn: any(
                                    c.name == "report"
                                    and c.ok
                                    and isinstance(c.output, dict)
                                    and c.output.get("verified")
                                    for c in turn.calls
                                )
                            )
                            if args.search
                            else None
                        ),
                    )
                except RuntimeError as exc:
                    print(f"  [mission] LLM error: {exc}")
                    break
                save_map()
                done = summary.get("success") if args.search else summary["interacted"]
                if done:
                    print(f"\n[mission] DONE: {summary['result']}")
                    return
                print(
                    f"  [mission] attempt {attempt} stopped ({summary['result']}, "
                    f"{summary['steps']} steps) — retrying"
                )
                controller.messages.append(
                    {
                        "role": "user",
                        "content": (
                            "You have not achieved the mission yet. "
                            + (
                                controller.state_hint() + " "
                                if controller.state_hint()
                                else ""
                            )
                            + "Keep trying: call tools and do not stop until the goal is achieved."
                        ),
                    }
                )
            print(
                f"[mission] gave up after {args.mission_retries + 1} attempts — "
                "say '/mission' to retry or command the capsule manually"
            )
        finally:
            saver.cancel()

    def start_mission(goal_text: str | None = None) -> None:
        nonlocal mission, mission_goal
        if mission is not None and not mission.done():
            print("[mission] already running — /stop first (or wait)")
            return
        mission_goal = goal_text or mission_goal
        mission = asyncio.create_task(mission_runner(mission_goal))

    async def stop_mission() -> None:
        nonlocal mission
        if mission is not None and not mission.done():
            mission.cancel()
            try:
                await mission
            except asyncio.CancelledError:
                pass
            mission = None
            print("[mission] stopped")
        else:
            print("  (no mission running)")

    async def cmd_status() -> None:
        if mission is not None and not mission.done():
            print(f"[mission] running: {mission_goal}")
        else:
            print("  (no mission running)")
        await cmd_state()

    if args.mission or args.search:
        if args.search:
            mission_goal = (
                "explore the room, find the orange triangle painted on the back "
                "of one of the crates, and report it"
            )
        start_mission()

    while True:
        # Surface a finished mission task so its result is reported once.
        if mission is not None and mission.done():
            try:
                mission.result()
            except Exception as exc:  # noqa: BLE001 - report mission failure
                print(f"[mission] error: {type(exc).__name__}: {exc}")
            mission = None
            print("[mission] finished — you can start another with /mission")
        try:
            # input() in a thread: a blocking read here would freeze the
            # event loop and stall a running mission task.
            text = (await asyncio.to_thread(input, "you> ")).strip()
        except (EOFError, KeyboardInterrupt):
            print("\n[chat] bye")
            return 0
        if not text:
            continue
        low = text.lower()
        if low in ("/exit", "exit", "quit", "выход"):
            if mission is not None and not mission.done():
                mission.cancel()
                try:
                    await mission
                except asyncio.CancelledError:
                    pass
            print("[chat] bye")
            return 0
        if low == "/help":
            print(HELP)
            continue
        if low == "/mission":
            start_mission()
            continue
        if low.startswith("/mission "):
            start_mission(low.split(maxsplit=1)[1])
            continue
        if low == "/stop":
            await stop_mission()
            continue
        if low == "/status":
            await cmd_status()
            continue
        if low == "/state":
            await cmd_state()
            continue
        if low == "/look":
            await cmd_look()
            continue
        if low == "/map":
            save_map()
            continue
        if low == "/models":
            await cmd_models()
            continue
        if low.startswith("/model "):
            model = low.split(maxsplit=1)[1]
            controller.model = model
            print(f"[chat] switching to {model}")
            continue
        if low.startswith("/steps "):
            try:
                auto_steps = max(0, int(low.split(maxsplit=1)[1]))
                print(f"[chat] auto-continue budget: {auto_steps} steps")
            except ValueError:
                print("  usage: /steps N")
            continue

        controller.messages.append({"role": "user", "content": text})
        steps_used = 0
        while True:
            try:
                turn = await controller.invoke()
            except RuntimeError as exc:
                print(f"  [error] {exc}")
                break
            steps_used += 1
            if turn.text:
                print(f"  model> {turn.text}")
            if turn.thinking:
                print(f"  [thinking] {turn.thinking[:300]}")
            if turn.interacted:
                print(f"\n[chat] BEACON ACTIVATED: {turn.message}")
                save_map()
                return 0
            if not turn.calls:
                if not turn.text and auto_steps and steps_used <= auto_steps:
                    # The model answered with nothing: nudge it to actually act.
                    controller.messages.append(
                        {
                            "role": "user",
                            "content": (
                                "Your last turn contained no action and no answer. "
                                "Carry out the user's request: call a tool now "
                                "(check the CURRENT STATE note)."
                            ),
                        }
                    )
                    print("  [chat] (nudge: model gave an empty turn)")
                    continue
                break  # the model answered in words; wait for the user
            if auto_steps and steps_used >= auto_steps:
                print(
                    f"  [chat] (auto-continue budget of {auto_steps} reached — say 'continue' to keep going)"
                )
                break
            if not auto_steps:
                break
            # The model acted without commenting; let it keep going for 'go to X'
        # Real-time perception in manual chat: refresh the model's view of the
        # world on a cadence, so it reacts to what it sees without being asked.
        user_turns += 1
        if args.look_every and user_turns % args.look_every == 0:
            await controller.auto_frame(
                "Fresh camera frame for your reference — react if the world changed."
            )
        if controller.pos is not None:
            save_map()

    return 0


async def run(args: argparse.Namespace) -> int:
    """Run the chat session, reconnecting to the bridge if the connection dies."""
    attempts = 0
    while True:
        session_rc: int | None = None
        try:
            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]}")
                session_rc = await chat_loop(client, manifest, args)
        except asyncio.CancelledError:
            raise
        except KeyboardInterrupt:
            raise
        except Exception as exc:  # noqa: BLE001 - connection lost: reconnect
            if session_rc is not None:
                # chat_loop exited cleanly; only cleanup failed — ignore.
                print(f"  (cleanup note: {type(exc).__name__})")
                return session_rc
            attempts += 1
            print(f"\n[chat] connection lost ({type(exc).__name__}: {exc})")
            if attempts >= 3:
                print(
                    "[chat] giving up after 3 attempts — is the bridge running? (scripts/run_bridge.sh)"
                )
                return 1
            print(
                "[chat] reconnecting in 2 seconds… (the room keeps its state on the bridge)"
            )
            await asyncio.sleep(2.0)
        else:
            return session_rc if session_rc is not None else 0


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Interactive chat: talk to the LLM, it drives the capsule."
    )
    parser.add_argument(
        "--url",
        default=DEFAULT_URL,
        help=f"bridge WebSocket URL (default {DEFAULT_URL})",
    )
    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(
        "--auto-steps",
        type=int,
        default=8,
        help="how many tool steps the model may chain per request (0 = one action per turn)",
    )
    parser.add_argument(
        "--mission",
        action="store_true",
        help="start an autonomous mission on connect: reach and activate the beacon",
    )
    parser.add_argument(
        "--search",
        action="store_true",
        help="search challenge: explore the room and find the orange triangle "
        "on the back of a crate, then report it (autonomous mode, fewer guardrails)",
    )
    parser.add_argument(
        "--free",
        action="store_true",
        help="fewer restrictions: no re-aim corrections or collision hints; the model plans freely",
    )
    parser.add_argument(
        "--mission-steps",
        type=int,
        default=50,
        help="tool steps per mission attempt (default 50)",
    )
    parser.add_argument(
        "--mission-retries",
        type=int,
        default=2,
        help="restarts after a failed attempt (default 2)",
    )
    parser.add_argument(
        "--mission-nudges",
        type=int,
        default=3,
        help="how many consecutive nudges a mission may use before giving up (default 3)",
    )
    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 turns (0 disables; default 3)",
    )
    parser.add_argument(
        "--cruise",
        type=float,
        default=1.0,
        help="proactive motion in missions: glide forward (meters) while the LLM thinks (0 disables)",
    )
    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"[chat] provider: {args.provider} -> {args.base_url}  model: {args.model}"
            )
    except ValueError as exc:
        print(f"[chat] {exc}")
        return 1
    try:
        return asyncio.run(run(args))
    except KeyboardInterrupt:
        print("\n[chat] bye")
        return 0


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