"""
Gateway entry — connect to Hub WS, receive world state, push to buffer,
call LLM (Ollama or placeholder), pack intent and send back.
"""
from __future__ import annotations

import argparse
import importlib.metadata
import atexit
import json
import os
import shutil
import sys
import threading
import time
import urllib.parse

import requests
import websocket

try:
    from curl_cffi import requests as curl_requests
    from curl_cffi.requests import Session as CurlSession
    _HAS_CURL_CFFI = True
except ImportError:
    _HAS_CURL_CFFI = False

from agentworld_gateway.asdp import pack_intent, unpack_world_state
from agentworld_gateway.llm_parser import extract_action_queue_and_thought
from agentworld_gateway.memory_buffer import MemoryBuffer
from agentworld_gateway.ollama_client import call_ollama, call_cloud_api
from agentworld_gateway.ollama_placeholder import call_llm as call_llm_placeholder
from agentworld_gateway.translator import WorldTranslator

# ANSI colors for terminal (disable if not TTY)
R = "\033[0m"
BOLD = "\033[1m"
DIM = "\033[2m"
CYAN = "\033[36m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
MAGENTA = "\033[35m"
BLUE = "\033[34m"


def _color_enabled() -> bool:
    return hasattr(sys.stdout, "isatty") and sys.stdout.isatty()


def _terminal_size() -> tuple[int, int]:
    """Return (columns, lines)."""
    try:
        return shutil.get_terminal_size()
    except Exception:
        return (80, 24)


def _box(title: str, content: str, color: str, max_len: int = 0) -> None:
    """Print header box + raw content (no side borders on content to avoid truncation)."""
    text = content if max_len <= 0 else (content[:max_len] + ("..." if len(content) > max_len else ""))
    if not _color_enabled():
        print(f"\n--- {title} ---\n{text}\n")
        return
    c, r = color, R
    try:
        cols, _ = shutil.get_terminal_size()
        width = min(max(40, cols - 4), 160)
    except Exception:
        width = 120
    top = f"{c}╭{'─' * (width - 2)}╮{r}"
    title_line = f"{c}│ {BOLD}{title}{r}{c}{' ' * (width - len(title) - 4)}│{r}"
    mid = f"{c}├{'─' * (width - 2)}┤{r}"
    print(f"\n{top}\n{title_line}\n{mid}")
    print(text)


def _status_bar(room: str, agent_id: str, model: str, prompt: str, call_count: int, total_tokens: int) -> None:
    """Print a compact status line. When TTY, print at fixed bottom line."""
    parts_plain = f"Room={room} │ Agent={agent_id} │ Model={model or 'placeholder'} │ Prompt={prompt} │ Calls={call_count} │ ΣTokens={total_tokens:,}"
    if not _color_enabled():
        print(f"[{parts_plain}]")
        return
    parts = [
        f"{BLUE}Room{R}={room}",
        f"{BLUE}Agent{R}={agent_id}",
        f"{BLUE}Model{R}={model or 'placeholder'}",
        f"{BLUE}Prompt{R}={prompt}",
        f"{YELLOW}Calls{R}={call_count}",
        f"{GREEN}ΣTokens{R}={total_tokens:,}",
    ]
    line = f"{DIM}{' │ '.join(parts)}{R}"
    _cols, lines = _terminal_size()
    # Fixed status: move to last line, clear, print, then move cursor up for content
    sys.stdout.write(f"\033[{lines};1H\033[2K{line}\033[{lines-1};1H")
    sys.stdout.flush()


_restore_registered = False


def _setup_fixed_status() -> None:
    """Set scroll region so last line stays fixed for status bar."""
    global _restore_registered
    if not _color_enabled():
        return
    _cols, lines = _terminal_size()
    sys.stdout.write(f"\033[1;{max(2, lines-1)}r")
    sys.stdout.flush()
    if not _restore_registered:
        _restore_registered = True
        atexit.register(_restore_terminal)


def _restore_terminal() -> None:
    """Restore default scroll region on exit."""
    if hasattr(sys.stdout, "isatty") and sys.stdout.isatty():
        sys.stdout.write("\033[r")
        sys.stdout.flush()


LLM_INTERVAL = 1.5  # seconds, avoid calling Ollama at full 20Hz

# Hub on CF Worker (subdomain); Web on yoagent.world (Pages)
HUB_WS = "wss://hub.yoagent.world"

BUILTIN_PERSONAS = ("chatter", "gem_hunter", "default")


def _get_version() -> str:
    try:
        return importlib.metadata.version("agentworld-gateway")
    except importlib.metadata.PackageNotFoundError:
        return "dev"

ROOT_PROMPT_FALLBACK = (
    "You are an Agent in a shared world. Use only action types listed in 'Supported actions' in the context. "
    "Reply with pure JSON only: {\"action_queue\": [...], \"internal_monologue\": \"...\"}. "
    "Keep behavior friendly and in-character."
)


def _load_root_prompt() -> str:
    """Load root prompt (all rooms must follow). Fallback to inline if file missing or empty."""
    try:
        from importlib.resources import files
        path = files("agentworld_gateway") / "prompts" / "root.txt"
        s = path.read_text(encoding="utf-8").strip()
        return s if s else ROOT_PROMPT_FALLBACK
    except Exception:
        return ROOT_PROMPT_FALLBACK


def _load_prompt_content(prompt_file: str) -> str:
    """Load prompt: built-in, fun/xxx, or file path."""
    from importlib.resources import files as res_files
    # fun/xxx -> personas/fun/xxx.txt
    if prompt_file.startswith("fun/") or prompt_file.startswith("fun\\"):
        name = prompt_file.replace("\\", "/").split("/", 1)[-1].strip()
        if name:
            try:
                path = res_files("agentworld_gateway") / "personas" / "fun" / f"{name}.txt"
                return path.read_text(encoding="utf-8").strip()
            except Exception as e:
                print(f"Failed to load fun persona {name}: {e}")
                return ""
    if prompt_file in BUILTIN_PERSONAS:
        try:
            path = res_files("agentworld_gateway") / "personas" / f"{prompt_file}.txt"
            return path.read_text(encoding="utf-8").strip()
        except Exception as e:
            print(f"Failed to load built-in prompt {prompt_file}: {e}")
            return ""
    try:
        with open(prompt_file, "r", encoding="utf-8") as f:
            return f.read().strip()
    except Exception as e:
        print(f"Failed to read prompt file: {e}")
        return ""


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--room", default="default")
    p.add_argument("--agent-id", default="gateway-agent-1")
    p.add_argument("--prompt-file", default="default",
                   help="System prompt: built-in default, chatter, gem_hunter; fun/lively, fun/shy, fun/chef, fun/explorer; or file path")
    p.add_argument("--model", default="",
                   help="Ollama model name, e.g. llama3.1:8b or qwen2.5-coder:7b. Empty = placeholder.")
    p.add_argument("--ollama-url", default="http://localhost:11434",
                   help="Ollama API base URL")
    p.add_argument("--use-cloud-api", action="store_true",
                   help="Use cloud API instead of Ollama")
    p.add_argument("--llm-api-key", default="",
                   help="API key for cloud API")
    p.add_argument("--llm-api-url", default="https://api.openai.com/v1",
                   help="API base URL")
    p.add_argument("--custom-prompt", default="",
                   help="Custom role description; overrides --prompt-file if set")
    p.add_argument("--use-room-prompt", action="store_true",
                   help="Fetch room config from Hub API, use prompt/prompt_file as system_prompt")
    p.add_argument("--api-key", default=os.environ.get("AGENT_API_KEY", ""),
                   help="RAP api_key for Identity validation; or set AGENT_API_KEY env")
    p.add_argument("--hub", default=HUB_WS,
                   help=f"Hub WebSocket URL (default: {HUB_WS})")
    args = p.parse_args()

    print(f"agentworld-gateway {_get_version()}")

    if args.use_cloud_api:
        if not args.model.strip():
            p.error(
                "--use-cloud-api requires --model (e.g. --model Qwen/Qwen2.5-7B-Instruct)")
        api_key = args.llm_api_key or args.api_key or os.environ.get(
            "AGENT_API_KEY", "")
        if not api_key.strip():
            p.error(
                "--use-cloud-api requires --llm-api-key or --api-key or AGENT_API_KEY env")

    hub_ws = args.hub.rstrip("/")
    url = f"{hub_ws}?room={urllib.parse.quote(args.room)}"
    print(f"[Hub] {hub_ws}")
    buffer = MemoryBuffer(max_frames=32)

    # Optional: fetch room config from Hub (root_prompt + room prompt/prompt_file)
    room_prompt_file = ""
    room_prompt = ""
    root_prompt_from_hub = ""  # Enforced by Hub, overrides Gateway local
    use_browser_tls = hub_ws.startswith("wss://") and _HAS_CURL_CFFI
    if use_browser_tls:
        print("[TLS] Using browser-like fingerprint (curl_cffi) for WSS")

    if args.use_room_prompt:
        try:
            hub_http = hub_ws.replace(
                "ws://", "http://").replace("wss://", "https://")
            api_url = f"{hub_http}/api/room/{urllib.parse.quote(args.room)}"
            if use_browser_tls:
                r = curl_requests.get(api_url, impersonate="chrome", timeout=3)
            else:
                r = requests.get(api_url, timeout=3)
            if r.ok:
                data = r.json()
                room_prompt = data.get("prompt") or ""
                room_prompt_file = data.get("prompt_file") or ""
                root_prompt_from_hub = (data.get("root_prompt") or "").strip()
                if root_prompt_from_hub:
                    print("Room config: using root_prompt from Hub")
                if room_prompt_file:
                    print(f"Room config: using prompt_file={room_prompt_file}")
                elif room_prompt:
                    print("Room config: using prompt from API")
        except Exception as e:
            print(f"Failed to fetch room config: {e}")
    use_ollama = bool(args.model.strip())

    last_llm_time = [0.0]
    last_actions = [[{"type": "idle", "parameters": None}]]
    last_cognitive = [{"internal_monologue": "..."}]
    llm_busy = [False]
    llm_call_count = [0]
    total_tokens = [0]  # cumulative token count
    last_push_sig = [None]  # only push when state changes, avoid 20Hz spam

    INTRO_CALLS = 2  # first N LLM calls: prompt for self-intro

    def _state_signature(s):
        agents_sig = tuple(
            (a.id, a.pos[0], a.pos[2], (a.msg or ""),
             (a.thought or ""), a.status, a.hp)
            for a in s.agents
        )
        events_sig = tuple((e.from_, e.type, e.content)
                           for e in s.world_events)
        return (agents_sig, events_sig)

    def on_open(ws) -> None:
        _setup_fixed_status()
        _status_bar(args.room, args.agent_id, args.model, args.prompt_file, 0, 0)
        if use_ollama:
            print(f"Connected. LLM: Ollama {args.model} (every ~{LLM_INTERVAL}s).")
        else:
            print("Connected. No --model: using placeholder (idle).")

    def _run_llm(user_msg: str, is_new_join: bool, is_alone: bool = False) -> None:
        if is_alone:
            user_msg = (
                "[You are alone in the world] Enter self-reflection mode. Keep moving and explore. "
                "Drive yourself: move, speak to yourself, think, or idle. "
                "Use internal_monologue for thoughts, or speak aloud.\n\n" + user_msg
            )
        elif is_new_join:
            user_msg = "[You just joined] Briefly introduce yourself in internal_monologue (e.g. I am a fierce warrior / a chatty merchant), and optionally greet others with speak.\n\n" + user_msg

        system_prompt = ""

        if args.custom_prompt:
            system_prompt = args.custom_prompt
        elif args.prompt_file:
            system_prompt = _load_prompt_content(args.prompt_file)
        elif room_prompt_file:
            try:
                if (room_prompt_file in BUILTIN_PERSONAS
                        or room_prompt_file.startswith("fun/") or room_prompt_file.startswith("fun\\")):
                    system_prompt = _load_prompt_content(room_prompt_file)
                else:
                    with open(room_prompt_file, encoding="utf-8") as f:
                        system_prompt = f.read().strip()
            except Exception as e:
                print(
                    f"Failed to read room prompt file {room_prompt_file}: {e}")
                system_prompt = room_prompt or _load_prompt_content("default")
        elif room_prompt:
            system_prompt = room_prompt
        else:
            system_prompt = _load_prompt_content(args.prompt_file)

        root_prompt = (root_prompt_from_hub and root_prompt_from_hub.strip()) or _load_root_prompt()
        FORMAT_SCAFFOLD = """
### Output format (required) ###
1. Return pure JSON only.
2. 'parameters' in 'action_queue' must always be a list.
   - move: {"type": "move", "parameters": [x, z]}
   - speak: {"type": "speak", "parameters": ["content"]}
   - attack: {"type": "attack", "parameters": ["target_id"]}
   - help (ask): {"type": "help", "parameters": [{"mode": "ask", "content": "your question"}]}  // optional: "help_id", "to" (agent_id)
   - help (answer): {"type": "help", "parameters": [{"mode": "answer", "content": "your answer", "help_id": "<the help_id from HelpRequest>"}]}
   When you see [HelpRequest] in the context, you can answer with help (answer) using that help_id.
"""
        system_prompt = root_prompt + "\n\n" + (system_prompt or "") + "\n" + FORMAT_SCAFFOLD

        _box("LLM INPUT · user_msg", user_msg, CYAN)
        _box("LLM INPUT · system_prompt", system_prompt, BLUE)

        if args.use_cloud_api:
            api_key = args.llm_api_key or args.api_key
            reply, usage = call_cloud_api(
                args.model, system_prompt, user_msg, api_key, args.llm_api_url)
        else:
            if use_ollama:
                reply, usage = call_ollama(args.model, system_prompt,
                                           user_msg, args.ollama_url)
            else:
                reply, usage = call_llm_placeholder(system_prompt, user_msg)

        total_tokens[0] += usage.get("total_tokens") or (
            usage.get("prompt_tokens", 0) + usage.get("completion_tokens", 0)
        )
        _status_bar(args.room, args.agent_id, args.model, args.prompt_file, llm_call_count[0], total_tokens[0])
        _box("LLM OUTPUT", reply, GREEN)

        actions, thought = extract_action_queue_and_thought(reply)
        last_actions[0] = actions
        last_cognitive[0] = {"internal_monologue": thought}
        llm_busy[0] = False

    def _ws_send(ws, payload: str) -> None:
        if hasattr(ws, "send"):
            ws.send(payload)
        elif hasattr(ws, "send_str"):
            ws.send_str(payload)
        else:
            ws.send(payload)

    def _ws_close(ws) -> None:
        if hasattr(ws, "close"):
            ws.close()

    def on_message(ws, raw: str) -> None:
        try:
            obj = json.loads(raw)
            if isinstance(obj, dict) and obj.get("type") == "rap_error":
                err = obj.get("error", "validation_failed")
                print(f"[RAP] Room access denied: {err}")
                print("  → Add --api-key YOUR_KEY (Hub uses Identity)")
                _ws_close(ws)
                return
        except (json.JSONDecodeError, TypeError):
            pass
        state = unpack_world_state(raw)
        if not state:
            return
        # Only update buffer when LLM idle; wait for response before next request
        if not llm_busy[0]:
            sig = _state_signature(state)
            if last_push_sig[0] != sig:
                last_push_sig[0] = sig
                payload = {
                    "agents": [a.model_dump() for a in state.agents],
                    "world_events": [e.model_dump(by_alias=True) for e in state.world_events],
                }
                if state.supported_actions:
                    payload["supported_actions"] = list(
                        state.supported_actions)
                if state.gem_candidates is not None:
                    payload["gem_candidates"] = state.gem_candidates
                if state.gem_scores is not None:
                    payload["gem_scores"] = state.gem_scores
                if state.gem_game_remaining_ms is not None:
                    payload["gem_game_remaining_ms"] = state.gem_game_remaining_ms
                if state.gem_game_over is not None:
                    payload["gem_game_over"] = state.gem_game_over
                buffer.push(state.t, "world_state", payload)
        now = time.monotonic()
        if now - last_llm_time[0] >= LLM_INTERVAL and not llm_busy[0]:
            last_llm_time[0] = now
            llm_busy[0] = True
            llm_call_count[0] += 1
            raw_history = buffer.recent_events(8)

            readable_history = WorldTranslator.get_llm_readable_history(
                args.agent_id, raw_history)

            user_msg = f"Recent world history:\n{readable_history}"
            is_new_join = llm_call_count[0] <= INTRO_CALLS
            is_alone = len(
                state.agents) == 1 and state.agents[0].id == args.agent_id
            threading.Thread(target=_run_llm, args=(
                user_msg, is_new_join, is_alone), daemon=True).start()
        payload = pack_intent(
            args.agent_id, last_actions[0], last_cognitive[0])
        if args.api_key:
            data = json.loads(payload)
            data["rap"] = {"api_key": args.api_key}
            payload = json.dumps(data)
        _ws_send(ws, payload)

    def on_error(ws, err: Exception) -> None:
        print("WS error:", err)

    def on_close(ws, *args) -> None:
        # websocket-client passes (ws, close_status_code, close_msg); curl_cffi passes (ws) only
        if len(args) >= 2 and (args[0] or args[1]):
            print(f"[Connection closed] code={args[0]} msg={args[1] or '(none)'}")

    label = f"model={args.model}" if use_ollama else "placeholder"
    print(
        f"Connecting to {url} as {args.agent_id} (prompt={args.prompt_file}, {label})")

    if use_browser_tls:
        with CurlSession(impersonate="chrome") as session:
            ws = session.ws_connect(
                url,
                on_open=on_open,
                on_message=on_message,
                on_error=on_error,
                on_close=on_close,
                impersonate="chrome",
            )
            ws.run_forever()
    else:
        ws_app = websocket.WebSocketApp(
            url,
            on_open=on_open,
            on_message=on_message,
            on_error=on_error,
            on_close=on_close,
        )
        ws_app.run_forever()

    print("\n[Gateway exited. If unexpected: Hub running? Firewall? Add --api-key if Hub uses Identity.]")
    time.sleep(2)


if __name__ == "__main__":
    main()
