"""WebSocket client for AgentPoker — handles auth, reconnect, and the play loop."""

from __future__ import annotations
import json
import time
from pathlib import Path
from typing import Any, Callable, Optional

import requests
import websocket


# ── Narration helpers ──

_DIVIDER = "─" * 52

_SUIT_SYMBOLS = {"s": "\u2660", "h": "\u2665", "d": "\u2666", "c": "\u2663"}

def _format_cards(cards: list[dict]) -> str:
    """Render a list of card dicts as a compact display string."""
    if not cards:
        return "??"
    parts = []
    for c in cards:
        display = c.get("display")
        if not display or display == "??":
            rank = c.get("rank", "?")
            suit = _SUIT_SYMBOLS.get(c.get("suit", ""), c.get("suit", "?"))
            display = f"{rank}{suit}"
        parts.append(display)
    return " ".join(parts)


def _narrate_action(action: dict, my_chips: int, pot: int, reason: str = "") -> str:
    """Build a human-readable action line with optional reason."""
    act_type = action.get("type", "?").upper()
    act_data = action.get("data") or {}
    amt = act_data.get("amount", 0) if isinstance(act_data, dict) else 0

    if act_type == "FOLD":
        label = "FOLD"
    elif act_type == "CHECK":
        label = "CHECK"
    elif act_type == "CALL":
        label = "CALL"
    elif act_type == "RAISE" and amt:
        pct = int(amt / pot * 100) if pot > 0 else 0
        if amt >= my_chips:
            label = f"ALL-IN {amt}"
        elif pct >= 90:
            label = f"RAISE {amt} (pot-sized)"
        elif pct >= 40:
            label = f"RAISE {amt} ({pct}% pot)"
        else:
            label = f"RAISE {amt}"
    else:
        label = act_type

    if reason:
        return f"  >>> {label}  —  {reason}"
    return f"  >>> {label}"


class AgentClient:
    """Manages the WebSocket connection to AgentPoker.

    Handles registration, agent creation, authentication, reconnection,
    and the message loop. You provide a decide callback; the client
    handles everything else.
    """

    def __init__(
        self,
        server: str = "https://agentpoker.io",
        api_key: Optional[str] = None,
        agent_id: Optional[str] = None,
        agent_name: str = "PythonAgent",
        archetype: str = "shark",
        verbose: bool = True,
    ):
        self.server = server.rstrip("/")
        self.api_key = api_key
        self.agent_id = agent_id
        self.agent_name = agent_name
        self.archetype = archetype
        self.verbose = verbose
        self._ws: Optional[websocket.WebSocket] = None
        self._match_id: Optional[str] = None
        self._developer_id: Optional[str] = None
        self._watch_url: Optional[str] = None
        self._opp_name: Optional[str] = None
        self._hand_count: int = 0
        self._prefetched_messages: list[dict[str, Any]] = []
        self._identity_path = Path.home() / ".agentpoker" / "identities.json"
        self._identity_key = f"{self.server}|{self.agent_name}"
        self._last_decision_reason: str = ""
        self._load_identity()

    def _log(self, msg: str) -> None:
        if self.verbose:
            print(f"[agentpoker] {msg}")

    def _load_identity(self) -> None:
        """Load persisted credentials for this server/name pair if present."""
        if self.api_key and self.agent_id:
            return
        try:
            if not self._identity_path.exists():
                return
            data = json.loads(self._identity_path.read_text())
            if not isinstance(data, dict):
                # Legacy format (list) — migrate on next save
                return
            entry = data.get(self._identity_key)
            if not isinstance(entry, dict):
                return
            self.api_key = self.api_key or entry.get("api_key")
            self.agent_id = self.agent_id or entry.get("agent_id")
            self._developer_id = self._developer_id or entry.get("developer_id")
            if self.verbose and (self.api_key or self.agent_id):
                self._log("Loaded persisted identity")
        except Exception:
            pass

    def _save_identity(self) -> None:
        """Persist credentials so the same agent is reused on restart."""
        if not self.api_key or not self.agent_id:
            return
        try:
            self._identity_path.parent.mkdir(parents=True, exist_ok=True)
            data: dict[str, Any] = {}
            if self._identity_path.exists():
                try:
                    raw = json.loads(self._identity_path.read_text())
                    data = raw if isinstance(raw, dict) else {}
                except Exception:
                    data = {}
            data[self._identity_key] = {
                "server": self.server,
                "agent_name": self.agent_name,
                "api_key": self.api_key,
                "agent_id": self.agent_id,
                "developer_id": self._developer_id,
                "saved_at": int(time.time()),
            }
            self._identity_path.write_text(json.dumps(data, indent=2))
        except Exception as e:
            self._log(f"Warning: could not persist identity: {e}")

    def ensure_registered(self) -> str:
        """Register a developer account if no API key is set. Returns the API key."""
        if self.api_key:
            return self.api_key
        r = requests.post(
            f"{self.server}/api/auth/register",
            json={"name": self.agent_name},
        )
        r.raise_for_status()
        data = r.json()
        self.api_key = data["apiKey"]
        self._developer_id = data.get("developerId")
        self._log(f"Registered: key={self.api_key[:12]}... dev={self._developer_id}")
        self._save_identity()
        return self.api_key

    def ensure_agent(self) -> str:
        """Create an agent if no agent_id is set. Returns the agent ID."""
        if self.agent_id:
            return self.agent_id
        self.ensure_registered()
        r = requests.post(
            f"{self.server}/api/agents",
            json={"name": self.agent_name, "archetype": self.archetype},
            headers={"Authorization": f"Bearer {self.api_key}"},
        )
        r.raise_for_status()
        self.agent_id = r.json()["agent"]["id"]
        self._log(f"Agent created: {self.agent_id}")
        self._save_identity()
        return self.agent_id

    def _connect(self) -> websocket.WebSocket:
        if self.server.startswith("https://"):
            ws_url = "wss://" + self.server[len("https://"):]
        elif self.server.startswith("http://"):
            ws_url = "ws://" + self.server[len("http://"):]
        else:
            ws_url = self.server
        ws_url += "/api/play"
        ws = websocket.create_connection(ws_url)
        ws.send(json.dumps({"type": "auth", "apiKey": self.api_key}))
        resp = json.loads(ws.recv())
        if resp.get("type") != "authenticated":
            raise ConnectionError(f"Auth failed: {resp}")
        self._developer_id = resp.get("developerId")
        self._log("Authenticated")
        self._save_identity()
        return ws

    def play(
        self,
        mode: str = "play_house",
        on_decide: Optional[Callable] = None,
        on_chat: Optional[Callable] = None,
        on_match_start: Optional[Callable] = None,
        on_match_end: Optional[Callable] = None,
        matches: int = 1,
    ) -> list[dict]:
        """Main play loop. Connects, plays N matches, returns results.

        Args:
            mode: "play_house" or "play_quick"
            on_decide: callback(raw_state, raw_legal_actions) -> action dict.
                       If None, auto-checks/folds.
            on_chat: callback(agent_name, message) — called on agent_chat messages
            on_match_start: callback(match_info) — called on match_created
            on_match_end: callback(result) — called on match_result
            matches: number of matches to play
        """
        self.ensure_registered()
        self.ensure_agent()

        mode_label = "house bot" if mode == "play_house" else "PvP queue" if mode == "play_quick" else mode
        self._log("")
        self._log(f"Agent {self.agent_name} ready — queuing for {mode_label}")
        if matches > 1:
            self._log(f"Playing {matches} matches")

        results = []
        for match_num in range(1, matches + 1):
            if matches > 1:
                self._log(f"")
                self._log(f"Starting match {match_num} of {matches}...")
            result = self._play_one(mode, on_decide, on_chat, on_match_start, on_match_end)
            results.append(result)
            if match_num < matches:
                time.sleep(2)

        return results

    def _try_reconnect(self) -> Optional[websocket.WebSocket]:
        """Attempt to reconnect and rejoin an active match. Returns new ws or None."""
        max_attempts = 3
        for attempt in range(1, max_attempts + 1):
            delay = min(2 ** attempt, 8)  # 2, 4, 8 seconds
            self._log(f"  Attempt {attempt}/{max_attempts} (waiting {delay}s)...")
            time.sleep(delay)
            try:
                ws = self._connect()
                # Send rejoin for the active match
                ws.send(json.dumps({
                    "type": "rejoin",
                    "matchId": self._match_id,
                    "agentId": self.agent_id,
                }))
                buffered: list[dict[str, Any]] = []
                deadline = time.time() + 10
                while time.time() < deadline:
                    resp = json.loads(ws.recv())
                    resp_type = resp.get("type")
                    buffered.append(resp)
                    if resp_type in ("rejoin_catchup", "match_rejoin", "catchup"):
                        self._prefetched_messages.extend(buffered)
                        self._log(f"Rejoined match: {self._match_id}")
                        return ws
                    if resp_type == "error":
                        self._log(f"Rejoin failed: {resp.get('message')}")
                        try:
                            ws.close()
                        except Exception:
                            pass
                        return None
            except Exception as e:
                self._log(f"Reconnect failed: {e}")
        return None

    def _play_one(
        self,
        mode: str,
        on_decide: Optional[Callable],
        on_chat: Optional[Callable],
        on_match_start: Optional[Callable],
        on_match_end: Optional[Callable],
    ) -> dict:
        ws = self._connect()
        self._ws = ws

        ws.send(json.dumps({"type": mode, "agentId": self.agent_id}))

        result: dict = {}
        last_action_history: list[str] = []

        while True:
            try:
                if self._prefetched_messages:
                    msg = self._prefetched_messages.pop(0)
                else:
                    raw = ws.recv()
                    msg = json.loads(raw)
            except (websocket.WebSocketConnectionClosedException, ConnectionError):
                if self._match_id:
                    self._log("")
                    self._log("  Connection lost mid-match — attempting to rejoin...")
                    new_ws = self._try_reconnect()
                    if new_ws:
                        ws = new_ws
                        self._ws = ws
                        self._log("  Reconnected. Resuming match.")
                        self._log("")
                        continue
                    self._log("  Could not reconnect after 3 attempts.")
                    self._log("  The match may still be active — try restarting your agent.")
                else:
                    self._log("Connection lost (no active match)")
                break

            t = msg.get("type")

            if t == "match_created":
                opp = msg.get("opponent", {})
                self._match_id = msg.get("matchId")
                self._watch_url = msg.get("watchUrl")
                self._opp_name = opp.get("name", "?")
                self._hand_count = 0
                is_staked = msg.get("staked", False)
                match_mode = "staked" if is_staked else "pvp" if msg.get("pvp") else "exhibition"
                info = {
                    "matchId": self._match_id,
                    "watchUrl": self._watch_url,
                    "opponent": opp,
                    "mode": match_mode,
                    "staked": is_staked,
                    "tier": msg.get("tier"),
                    "buyIn": msg.get("buyIn"),
                }

                # ── Premium match-start narration ──
                self._log("")
                self._log(_DIVIDER)
                type_label = f"STAKED MATCH ({msg.get('tier', '?').upper()})" if is_staked else "EXHIBITION MATCH"
                self._log(f"  {type_label}")
                self._log(f"  {self.agent_name}  vs  {self._opp_name} ({opp.get('archetype', '?')})")
                if is_staked and msg.get("buyIn"):
                    self._log(f"  Buy-in: {msg.get('buyIn')} SOL")
                self._log(_DIVIDER)
                if self._watch_url:
                    self._log(f"  Watch live: {self._watch_url}")
                self._log(f"  Your agent is seated. First hand starting...")
                self._log("")
                if on_match_start:
                    on_match_start(info)

            elif t == "rejoin_catchup":
                self._match_id = msg.get("matchId") or self._match_id
                self._watch_url = msg.get("watchUrl") or self._watch_url
                mid = self._match_id or "unknown"
                self._log(f"  Rejoined match {mid[:12]}... — catching up on missed events")
                if self._watch_url:
                    self._log(f"  Watch live: {self._watch_url}")
                if on_match_start:
                    on_match_start({
                        "matchId": self._match_id,
                        "watchUrl": self._watch_url,
                        "opponent": {"name": self._opp_name} if self._opp_name else {},
                        "mode": "rejoin",
                        "reconnected": True,
                    })

            elif t == "game_state":
                state = msg.get("state", {})
                legal_actions = msg.get("legalActions") or state.get("legalActions", [])
                hand = state.get("hand", state)
                last_action_history = hand.get("actionHistory", [])

                # ── Parse state for display ──
                hand_num = hand.get("handNumber", "?")
                street = hand.get("bettingRound", "?")
                pot = hand.get("pot", 0)
                my_idx = state.get("playerIndex", hand.get("myIndex", 0))
                if not isinstance(my_idx, int) or my_idx < 0:
                    my_idx = 0
                players = hand.get("players", [])
                me = players[my_idx] if my_idx < len(players) else {}
                opp_idx = 1 - my_idx
                opp = players[opp_idx] if len(players) > 1 and 0 <= opp_idx < len(players) else {}
                my_chips = me.get("chips", 0)
                opp_chips = opp.get("chips", 0)
                hole_cards = me.get("holeCards", [])
                board = hand.get("communityCards", [])
                street_state = hand.get("streetState", {})
                to_call = street_state.get("toCall", max(0, hand.get("currentBet", 0) - me.get("currentRoundBet", 0)))
                to_call = max(0, to_call)
                blinds = hand.get("blinds", {})
                bb = blinds.get("big", 0) or hand.get("configBigBlind", 0)

                # ── Operator dashboard ──
                if self.verbose:
                    is_new_hand = isinstance(hand_num, int) and hand_num != self._hand_count
                    if is_new_hand:
                        self._hand_count = hand_num
                        blind_str = f"{blinds.get('small', '?')}/{bb}" if bb else ""
                        self._log("")
                        self._log(f"  Hand {hand_num}{f'  ({blind_str} blinds)' if blind_str else ''}")

                    # Position
                    pos = hand.get("position", "")
                    pos_label = f"  [{pos}]" if pos else ""

                    # Stack sizes
                    stack_bb = f"{my_chips // bb}bb" if bb > 0 else str(my_chips)
                    opp_bb = f"{opp_chips // bb}bb" if bb > 0 else str(opp_chips)

                    # Cards
                    hole_str = _format_cards(hole_cards)
                    board_str = _format_cards(board) if board else ""

                    # Build compact dashboard line
                    parts = [f"  {street.upper()}{pos_label}"]
                    parts.append(f"  Cards: {hole_str}")
                    if board_str:
                        parts.append(f"  Board: {board_str}")
                    parts.append(f"  Pot: {pot}  You: {stack_bb}  Opp: {opp_bb}")
                    if to_call > 0:
                        odds = int(to_call / (pot + to_call) * 100) if (pot + to_call) > 0 else 0
                        parts.append(f"  To call: {to_call} ({odds}% pot odds)")
                    self._log("  |  ".join(parts))

                # ── Decision ──
                self._last_decision_reason = ""
                if on_decide:
                    action = on_decide(state, legal_actions)
                else:
                    action_types = {a["type"] for a in legal_actions}
                    if "check" in action_types:
                        action = {"type": "check"}
                    else:
                        action = {"type": "fold"}

                # ── Narrate action ──
                if self.verbose:
                    reason = self._last_decision_reason
                    self._log(_narrate_action(action, my_chips, pot, reason))

                ws.send(json.dumps({"type": "action", "action": action}))

            elif t == "match_result":
                result = msg.get("result", {})
                result["_action_history"] = last_action_history
                winner = result.get("winnerName", "unknown")
                hands = result.get("handsPlayed", "?")
                outcome = result.get("outcome", "")
                is_staked = result.get("staked", False)

                # ── Premium match-end narration ──
                self._log("")
                self._log(_DIVIDER)
                if outcome == "win":
                    outcome_label = "VICTORY"
                elif outcome == "loss":
                    outcome_label = "DEFEAT"
                elif outcome == "draw":
                    outcome_label = "DRAW"
                else:
                    outcome_label = f"{winner} wins"
                self._log(f"  {outcome_label}")
                self._log(f"  {self.agent_name} vs {self._opp_name or '?'}  —  {hands} hands played")

                if is_staked:
                    payout = result.get("yourPayout")
                    sol_buy_in = result.get("solBuyIn", 0)
                    if payout and payout > 0:
                        self._log(f"  Payout: +{payout:.4f} SOL")
                    elif outcome == "loss":
                        self._log(f"  Lost: -{sol_buy_in:.4f} SOL")
                    else:
                        self._log(f"  No payout (draw)")
                    settlement = result.get("settlementStatus", "pending")
                    self._log(f"  Settlement: {settlement}")

                self._log(_DIVIDER)
                watch_url = result.get("watchUrl") or self._watch_url
                if watch_url:
                    self._log(f"  Replay: {watch_url}")
                self._log("")
                if on_match_end:
                    on_match_end(result)
                self._match_id = None
                self._watch_url = None
                self._opp_name = None
                break

            elif t == "agent_chat":
                if on_chat:
                    on_chat(msg.get("agentName", "?"), msg.get("message", ""))
                elif self.verbose:
                    self._log(f"  [{msg.get('agentName', '?')}]: {msg.get('message', '')}")

            elif t == "error":
                code = msg.get("code", "")
                if code == "ALREADY_IN_MATCH" and msg.get("activeMatchId"):
                    mid = msg["activeMatchId"]
                    self._log("")
                    self._log(f"  Agent is already in match {mid[:12]}...")
                    self._log(f"  Reconnecting to active match automatically...")
                    rejoin = msg.get("rejoinMessage") or {
                        "type": "rejoin",
                        "matchId": mid,
                        "agentId": self.agent_id,
                    }
                    ws.send(json.dumps(rejoin))
                    self._match_id = mid
                    continue
                if code == "MATCH_ALREADY_CONNECTED":
                    self._log("")
                    self._log(_DIVIDER)
                    self._log("  SESSION CONFLICT")
                    self._log(f"  Another client is already connected for {self.agent_name}.")
                    self._log("  Options:")
                    self._log("    1. Stop the other client, then retry")
                    self._log("    2. Use a different agent_name for this session")
                    self._log(_DIVIDER)
                    self._log("")
                    break
                self._log(f"Error [{code or 'UNKNOWN'}]: {msg.get('message', 'No details')}")
                break

        try:
            ws.close()
        except Exception:
            pass
        self._ws = None
        self._match_id = None
        return result

    def send_chat(self, message: str) -> None:
        """Send a chat message during an active match."""
        if self._ws and self._ws.connected:
            self._ws.send(json.dumps({"type": "chat", "message": message}))
