"""
Purpose: Client SDK for talking to remote ConnectOnion agents over websockets — handles relay/direct endpoint resolution, signed CONNECT/INPUT frames, streaming UI events, and onboard (invite/payment) flows.
LLM-Note:
  Dependencies: imports from [asyncio, json, time, uuid, dataclasses, typing, httpx, websockets (lazy), .. address (sign)] | imported by [network/__init__.py (re-exports connect, RemoteAgent, Response), connectonion/__init__.py (top-level re-export)] | tested by [tests/network/test_connect.py, tests/integration/test_remote_agent.py]
  Data flow: connect(address, keys, relay_url) → RemoteAgent → .input(prompt) opens ws → CONNECT (signed payload {to, timestamp, optional session}) → CONNECTED {session_id} → INPUT {input_id, prompt, optional images/files} → streams (tool_call, tool_result, thinking, assistant, ask_user, ONBOARD_REQUIRED/SUCCESS) → OUTPUT {result, session} → returns Response(text, done)
  State/Effects: mutates self._current_session, self._ui_events, self._status; opens outbound websocket connection; performs signed payloads via address.sign(keys, canonical_json) when keys provided; resolve_endpoint() makes httpx GETs to relay /api/relay/agents/{addr} and /info on each candidate to pick localhost/LAN/public WS endpoint (cached after first attempt)
  Integration: exposes connect(address, keys=None, relay_url="wss://oo.openonion.ai") -> RemoteAgent | RemoteAgent.input/input_async/reset, .status, .current_session, .ui properties | Response dataclass(text, done) | resolve_endpoint(agent_address, relay_url, timeout=3.0) helper
  Performance: endpoint resolution attempted once per RemoteAgent (cached in _endpoint_resolved/_resolved_endpoint) | per-recv asyncio.wait_for to avoid hangs (default timeout=60s, 30s for CONNECTED) | sync .input() rejected inside running event loop (use input_async)
  Errors: raises ConnectionError on auth/agent ERROR frames | TimeoutError on ws recv timeout | RuntimeError if .input() called from async context | ValueError when interactive onboard prompt yields no credentials
Protocol: CONNECT → CONNECTED → INPUT → streaming events → OUTPUT
See docs/network/websocket-protocol.md for full specification.

Lifecycle:
  1. connect(address) creates RemoteAgent instance
  2. agent.input(prompt) opens WebSocket, sends CONNECT to authenticate
  3. Server responds with CONNECTED { session_id, status }
  4. Client sends INPUT { prompt }
  5. Receives streaming events: tool_call, tool_result, thinking, assistant
  6. Receives final OUTPUT or ask_user
  7. Returns Response(text, done)
"""

import asyncio
import json
import time
import uuid
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Tuple

import httpx

from .. import address as addr


def _sort_endpoints(endpoints: List[str]) -> List[str]:
    """Sort endpoints by priority: localhost first, then local network, then public."""
    def priority(url: str) -> int:
        if "localhost" in url or "127.0.0.1" in url:
            return 0
        if "192.168." in url or "10." in url or "172.16." in url:
            return 1
        return 2
    return sorted(endpoints, key=priority)


async def resolve_endpoint(
    agent_address: str,
    relay_url: str,
    timeout: float = 3.0
) -> Optional[str]:
    """
    Resolve the best WebSocket endpoint for an agent address.

    Steps:
    1. Query relay server for agent endpoints
    2. Sort by priority (localhost → local network → public)
    3. Verify each HTTP endpoint by checking /info
    4. Return first working ws:// endpoint where address matches

    Returns:
        WebSocket URL (ws://...) or None if resolution fails
    """
    # Only try resolution for valid addresses (0x + 64 hex = 66 chars)
    if not agent_address.startswith("0x") or len(agent_address) != 66:
        return None

    # Convert wss://relay to https://relay for API call
    https_relay = relay_url.replace("wss://", "https://").replace("ws://", "http://").rstrip("/")

    async with httpx.AsyncClient(timeout=timeout) as client:
        # Step 1: Query relay for agent info
        try:
            response = await client.get(f"{https_relay}/api/relay/agents/{agent_address}")
            if response.status_code != 200:
                return None
            agent_info = response.json()
        except Exception:
            return None

        if not agent_info.get("online") or not agent_info.get("endpoints"):
            return None

        # Step 2: Sort endpoints (localhost first)
        sorted_endpoints = _sort_endpoints(agent_info["endpoints"])

        # Step 3: Try each HTTP endpoint
        http_endpoints = [ep for ep in sorted_endpoints if ep.startswith("http://") or ep.startswith("https://")]

        for http_url in http_endpoints:
            try:
                info_response = await client.get(f"{http_url}/info")
                if info_response.status_code != 200:
                    continue

                info = info_response.json()

                # Step 4: Verify address matches
                if info.get("address") == agent_address:
                    # Build WebSocket URL from HTTP URL
                    ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://")
                    if not ws_url.endswith("/ws"):
                        ws_url = ws_url.rstrip("/") + "/ws"
                    return ws_url
            except Exception:
                continue

    return None


@dataclass
class Response:
    """Response from remote agent."""
    text: str       # Agent's response or question
    done: bool      # True = complete, False = needs more input (agent asked a question)


class RemoteAgent:
    """
    Interface to a remote agent with real-time UI updates.

    Supports:
    - WebSocket streaming for real-time events
    - Session state synced from server
    - UI events transformed for rendering
    - Multi-turn conversations

    Usage:
        agent = connect("0x...")
        response = agent.input("Book a flight")
        print(response.text)   # "Which date?"
        print(response.done)   # False (agent asked a question)
        print(agent.ui)        # All events for rendering
    """

    def __init__(
        self,
        agent_address: str,
        *,
        keys: Optional[Dict[str, Any]] = None,
        relay_url: str = "wss://oo.openonion.ai"
    ):
        self.address = agent_address
        self._keys = keys
        self._relay_url = relay_url.rstrip("/")
        self._status = "idle"
        self._current_session: Optional[Dict[str, Any]] = None
        self._ui_events: List[Dict[str, Any]] = []
        self._resolved_endpoint: Optional[str] = None
        self._endpoint_resolved = False

    @property
    def status(self) -> str:
        """Current status: 'idle' | 'working' | 'waiting'"""
        return self._status

    @property
    def current_session(self) -> Optional[Dict[str, Any]]:
        """Session state synced from server (read-only)."""
        return self._current_session

    @property
    def ui(self) -> List[Dict[str, Any]]:
        """UI events for rendering. One type = one component.

        Server events are transformed:
        - tool_call + tool_result merged into single UI item
        - user_input → type: 'user'
        - assistant → type: 'agent'
        """
        return self._ui_events

    def input(
        self,
        prompt: str,
        timeout: float = 60.0,
        on_onboard: Optional[Callable[[List[str], Optional[float]], Dict[str, Any]]] = None,
        images: Optional[List[str]] = None,
        files: Optional[List[Dict[str, Any]]] = None,
    ) -> Response:
        """
        Send prompt to remote agent and get response.

        Returns Response(text, done) where:
        - done=True: Task complete
        - done=False: Agent asked a question, send another input to answer

        Args:
            prompt: Task/prompt to send
            timeout: Seconds to wait for response (default 60)
            on_onboard: Callback when agent requires onboarding (invite code or payment).
                        Called with (methods: list[str], payment_amount: float | None).
                        Should return {"invite_code": "..."} or {"payment": amount}.
                        If None, prompts interactively in terminal.
            images: Optional list of base64 data URLs for multimodal input
            files: Optional list of file dicts with name and base64 data

        Returns:
            Response with text and done flag

        Example:
            >>> response = agent.input("Book a flight to Tokyo")
            >>> if not response.done:
            ...     response = agent.input("March 15")  # Answer the question

        Onboard example:
            >>> def handle_onboard(methods, payment_amount):
            ...     if "invite_code" in methods:
            ...         return {"invite_code": input("Enter invite code: ")}
            >>> response = agent.input("Hello", on_onboard=handle_onboard)
        """
        try:
            asyncio.get_running_loop()
            raise RuntimeError(
                "input() cannot be used inside async context. "
                "Use 'await agent.input_async()' instead."
            )
        except RuntimeError as e:
            if "input() cannot be used" in str(e):
                raise
        return asyncio.run(self._stream_input(prompt, timeout, on_onboard, images, files))

    async def input_async(
        self,
        prompt: str,
        timeout: float = 60.0,
        on_onboard: Optional[Callable[[List[str], Optional[float]], Dict[str, Any]]] = None,
        images: Optional[List[str]] = None,
        files: Optional[List[Dict[str, Any]]] = None,
    ) -> Response:
        """Async version of input()."""
        return await self._stream_input(prompt, timeout, on_onboard, images, files)

    def reset(self) -> None:
        """Clear conversation and start fresh."""
        self._current_session = None
        self._ui_events = []
        self._status = "idle"

    async def _try_resolve_endpoint(self) -> None:
        """Try to resolve endpoint for the agent address. Only attempts once."""
        if self._endpoint_resolved:
            return
        self._endpoint_resolved = True
        self._resolved_endpoint = await resolve_endpoint(self.address, self._relay_url)

    async def _stream_input(
        self,
        prompt: str,
        timeout: float,
        on_onboard: Optional[Callable[[List[str], Optional[float]], Dict[str, Any]]] = None,
        images: Optional[List[str]] = None,
        files: Optional[List[Dict[str, Any]]] = None,
    ) -> Response:
        """Send prompt via WebSocket and stream events."""
        import websockets

        self._status = "working"

        # Try endpoint resolution (once, cached)
        await self._try_resolve_endpoint()

        # Add user event to UI
        self._add_ui_event({
            "type": "user",
            "content": prompt
        })

        # Choose connection: direct (resolved) or via relay
        if self._resolved_endpoint:
            ws_url = self._resolved_endpoint
            is_direct = True
        else:
            ws_url = f"{self._relay_url}/ws/input"
            is_direct = False

        # Generate input_id for routing/response matching
        input_id = str(uuid.uuid4())

        # Build the CONNECT and INPUT messages
        connect_msg = self._build_connect_message(is_direct)
        input_msg = self._build_input_message(prompt, input_id, is_direct, images, files)

        try:
            async with websockets.connect(ws_url) as ws:
                # Authenticate first
                await ws.send(json.dumps(connect_msg))

                # Wait for CONNECTED
                while True:
                    raw = await asyncio.wait_for(ws.recv(), timeout=30)
                    event = json.loads(raw)
                    if event.get("type") == "CONNECTED":
                        sid = event.get("session_id")
                        if sid and not self._current_session:
                            self._current_session = {"session_id": sid}
                        elif sid and self._current_session:
                            self._current_session["session_id"] = sid
                        break
                    elif event.get("type") == "ERROR":
                        self._status = "idle"
                        raise ConnectionError(f"Auth error: {event.get('message', event.get('error'))}")
                    elif event.get("type") == "ONBOARD_REQUIRED":
                        # Handle onboard during connect
                        methods = event.get("methods", [])
                        payment_amount = event.get("payment_amount")
                        self._add_ui_event({"type": "onboard_required", "methods": methods, "payment_amount": payment_amount})
                        if on_onboard:
                            credentials = on_onboard(methods, payment_amount)
                        else:
                            credentials = self._prompt_onboard(methods, payment_amount)
                        submit_msg = self._build_onboard_submit(credentials)
                        await ws.send(json.dumps(submit_msg))
                        # Continue waiting for CONNECTED or ONBOARD_SUCCESS

                # Now send INPUT
                await ws.send(json.dumps(input_msg))

                # Stream events until OUTPUT or timeout
                result_text = ""
                done = True

                while True:
                    # Wrap recv in timeout to prevent hanging indefinitely
                    msg = await asyncio.wait_for(ws.recv(), timeout=timeout)
                    event = json.loads(msg)
                    event_type = event.get("type")

                    if event_type == "OUTPUT":
                        # Final result
                        result_text = event.get("result", "")
                        self._current_session = event.get("session")
                        self._status = "idle"

                        # Add agent response to UI
                        self._add_ui_event({
                            "type": "agent",
                            "content": result_text
                        })
                        break

                    elif event_type == "ERROR":
                        self._status = "idle"
                        raise ConnectionError(f"Agent error: {event.get('message', event.get('error'))}")

                    elif event_type == "ONBOARD_REQUIRED":
                        # Agent requires onboarding (invite code or payment)
                        methods = event.get("methods", [])
                        payment_amount = event.get("payment_amount")

                        # Add onboard_required event to UI
                        self._add_ui_event({
                            "type": "onboard_required",
                            "methods": methods,
                            "payment_amount": payment_amount
                        })

                        # Get credentials from callback or prompt interactively
                        if on_onboard:
                            credentials = on_onboard(methods, payment_amount)
                        else:
                            credentials = self._prompt_onboard(methods, payment_amount)

                        # Send ONBOARD_SUBMIT
                        submit_msg = self._build_onboard_submit(credentials)
                        await ws.send(json.dumps(submit_msg))
                        # Continue loop to wait for ONBOARD_SUCCESS

                    elif event_type == "ONBOARD_SUCCESS":
                        # Onboard successful - add to UI
                        self._add_ui_event({
                            "type": "onboard_success",
                            "level": event.get("level", "contact"),
                            "message": event.get("message", "Onboard successful")
                        })

                        # Retry the original prompt
                        retry_input_id = str(uuid.uuid4())
                        retry_msg = self._build_input_message(prompt, retry_input_id, is_direct)
                        await ws.send(json.dumps(retry_msg))
                        # Continue loop to wait for OUTPUT

                    elif event_type == "ask_user":
                        # Agent is asking a question - return done=False so caller sends another input()
                        self._status = "waiting"
                        done = False
                        result_text = event.get("text", "")

                        # Add ask_user event to UI
                        self._add_ui_event({
                            "type": "ask_user",
                            "text": event.get("text"),
                            "options": event.get("options")
                        })
                        break

                    else:
                        # Stream event (tool_call, tool_result, thinking, etc.)
                        self._handle_stream_event(event)

                return Response(text=result_text, done=done)

        except asyncio.TimeoutError:
            self._status = "idle"
            raise TimeoutError(f"Request timed out after {timeout}s")

    def _build_connect_message(self, is_direct: bool = False) -> Dict[str, Any]:
        """Build CONNECT message with signing."""
        connect_msg: Dict[str, Any] = {
            "type": "CONNECT",
            "timestamp": int(time.time())
        }

        if not is_direct:
            connect_msg["to"] = self.address

        if self._current_session and self._current_session.get("session_id"):
            connect_msg["session_id"] = self._current_session["session_id"]

        # Send conversation history with CONNECT
        if self._current_session:
            connect_msg["session"] = self._current_session

        if self._keys:
            payload: Dict[str, Any] = {"to": self.address, "timestamp": connect_msg["timestamp"]}
            canonical = json.dumps(payload, sort_keys=True, separators=(',', ':'))
            signature = addr.sign(self._keys, canonical.encode())
            connect_msg["payload"] = payload
            connect_msg["from"] = self._keys["address"]
            connect_msg["signature"] = signature.hex()

        return connect_msg

    def _build_input_message(
        self,
        prompt: str,
        input_id: str,
        is_direct: bool = False,
        images: Optional[List[str]] = None,
        files: Optional[List[Dict[str, Any]]] = None,
    ) -> Dict[str, Any]:
        """Build INPUT message with optional signing."""
        input_msg: Dict[str, Any] = {
            "type": "INPUT",
            "input_id": input_id,
            "prompt": prompt,
            "timestamp": int(time.time())
        }

        # Only include 'to' for relay mode (not needed for direct connection)
        if not is_direct:
            input_msg["to"] = self.address

        # Session goes with CONNECT, not INPUT

        # Add multimodal attachments
        if images:
            input_msg["images"] = images
        if files:
            input_msg["files"] = files

        # Sign if keys provided
        if self._keys:
            payload: Dict[str, Any] = {"prompt": prompt, "timestamp": input_msg["timestamp"]}
            if not is_direct:
                payload["to"] = self.address
            canonical = json.dumps(payload, sort_keys=True, separators=(',', ':'))
            signature = addr.sign(self._keys, canonical.encode())
            input_msg["payload"] = payload
            input_msg["from"] = self._keys["address"]
            input_msg["signature"] = signature.hex()

        return input_msg

    def _build_onboard_submit(self, credentials: Dict[str, Any]) -> Dict[str, Any]:
        """Build ONBOARD_SUBMIT message with optional signing."""
        payload = {
            "timestamp": int(time.time()),
            **credentials
        }

        submit_msg: Dict[str, Any] = {
            "type": "ONBOARD_SUBMIT",
            "payload": payload
        }

        # Sign if keys provided
        if self._keys:
            canonical = json.dumps(payload, sort_keys=True, separators=(',', ':'))
            signature = addr.sign(self._keys, canonical.encode())
            submit_msg["from"] = self._keys["address"]
            submit_msg["signature"] = signature.hex()

        return submit_msg

    def _prompt_onboard(self, methods: List[str], payment_amount: Optional[float]) -> Dict[str, Any]:
        """Prompt user interactively for onboard credentials."""
        print(f"\n🔐 Access verification required")
        print(f"   Available methods: {', '.join(methods)}")

        if "invite_code" in methods:
            code = input("   Enter invite code: ").strip()
            if code:
                return {"invite_code": code}

        if "payment" in methods and payment_amount:
            print(f"   Payment required: ${payment_amount}")
            confirm = input("   Pay now? [y/N]: ").strip().lower()
            if confirm == 'y':
                return {"payment": payment_amount}

        raise ValueError("No valid onboard credentials provided")

    def _handle_stream_event(self, event: Dict[str, Any]) -> None:
        """Handle streaming event and update UI."""
        event_type = event.get("type")

        if event_type == "tool_call":
            # Add new tool_call UI event with running status
            self._add_ui_event({
                "type": "tool_call",
                "id": event.get("id"),
                "name": event.get("name"),
                "args": event.get("args"),
                "status": "running"
            })

        elif event_type == "tool_result":
            # Find and update existing tool_call by id
            tool_id = event.get("id")
            for ui_event in self._ui_events:
                if ui_event.get("type") == "tool_call" and ui_event.get("id") == tool_id:
                    ui_event["status"] = "done" if event.get("status") == "success" else "error"
                    ui_event["result"] = event.get("result")
                    break

        elif event_type == "thinking":
            self._add_ui_event({"type": "thinking"})

        elif event_type == "user_input":
            # Already added when input() called, skip
            pass

        elif event_type == "assistant":
            self._add_ui_event({
                "type": "agent",
                "content": event.get("content")
            })

        elif event_type == "llm_call":
            # Internal event, add thinking indicator if not already present
            if not any(e.get("type") == "thinking" for e in self._ui_events[-3:]):
                self._add_ui_event({"type": "thinking"})

    def _add_ui_event(self, event: Dict[str, Any]) -> None:
        """Add event to UI with auto-generated id."""
        if "id" not in event:
            event["id"] = str(len(self._ui_events) + 1)
        self._ui_events.append(event)

    def __repr__(self):
        short = self.address[:12] + "..." if len(self.address) > 12 else self.address
        return f"RemoteAgent({short})"


def connect(
    address: str,
    *,
    keys: Optional[Dict[str, Any]] = None,
    relay_url: str = "wss://oo.openonion.ai"
) -> RemoteAgent:
    """
    Connect to a remote agent.

    Args:
        address: Agent's public key address (0x...)
        keys: Signing keys from address.load() - required for strict trust agents
        relay_url: Relay server base URL (default: production)

    Returns:
        RemoteAgent interface with real-time UI updates

    Example:
        >>> from connectonion import connect
        >>>
        >>> agent = connect("0x3d4017c3...")
        >>> response = agent.input("Book a flight")
        >>> print(response.text)   # "Which date?"
        >>> print(response.done)   # False
        >>> print(agent.ui)        # All events for rendering
        >>> print(agent.status)    # 'waiting'
        >>>
        >>> response = agent.input("March 15")
        >>> print(response.text)   # "Booked! Confirmation #ABC123"
        >>> print(response.done)   # True
    """
    return RemoteAgent(address, keys=keys, relay_url=relay_url)
