"""WebSocket chat client for interactive conversations with Navari agent."""

import asyncio
import json
import ssl
import traceback
from collections.abc import Callable, Coroutine
from pathlib import Path
from typing import Any

import websockets
from rich.console import Console
from rich.markdown import Markdown
from rich.panel import Panel

from navari_cli.api_client import NavariAPIClient

console = Console()


def _build_ws_url(base_url: str, conversation_id: str, token: str) -> str:
    """Build WebSocket URL for starting a new run."""
    ws_url = base_url.replace("https://", "wss://").replace("http://", "ws://")
    return f"{ws_url}/api/ws/conversations/{conversation_id}/start?token={token}"


def _build_ws_attach_url(base_url: str, conversation_id: str, token: str) -> str:
    """Build WebSocket URL for attaching to an existing run."""
    ws_url = base_url.replace("https://", "wss://").replace("http://", "ws://")
    return f"{ws_url}/api/ws/conversations/{conversation_id}/attach?token={token}"


def _build_ssl_context(verify_ssl: bool) -> ssl.SSLContext | None:
    """Build SSL context for WebSocket connection."""
    if not verify_ssl:
        ctx = ssl.create_default_context()
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE
        return ctx
    return None


def _extract_json_string_value(partial_json: str, key: str) -> str | None:
    """Extract a string value from a possibly incomplete JSON object.

    Given partial JSON like '{"message": "Hello wor' and key 'message',
    returns 'Hello wor'. Returns None if the key hasn't appeared yet.

    Handles JSON escape sequences (\\n, \\", etc.) by decoding them.
    """
    # Look for "key": " or "key":" patterns
    for pattern in [f'"{key}": "', f'"{key}":"']:
        idx = partial_json.find(pattern)
        if idx >= 0:
            start = idx + len(pattern)
            # Extract everything after the opening quote, handling escapes
            result = []
            i = start
            while i < len(partial_json):
                ch = partial_json[i]
                if ch == "\\" and i + 1 < len(partial_json):
                    # Escaped character
                    next_ch = partial_json[i + 1]
                    if next_ch == "n":
                        result.append("\n")
                    elif next_ch == "t":
                        result.append("\t")
                    elif next_ch == '"':
                        result.append('"')
                    elif next_ch == "\\":
                        result.append("\\")
                    else:
                        result.append(next_ch)
                    i += 2
                elif ch == '"':
                    break  # End of string value
                else:
                    result.append(ch)
                    i += 1
            return "".join(result)
    return None


def _print_todo_summary(content: str) -> None:
    """Print a compact summary of a todo list from write_to_do_file content."""
    lines = content.strip().splitlines()
    done = 0
    pending = 0
    items = []

    for line in lines:
        stripped = line.strip()
        if stripped.startswith("- [x]"):
            done += 1
            items.append(f"  [green]✓[/green] [dim]{stripped[6:].strip()}[/dim]")
        elif stripped.startswith("- [ ]"):
            pending += 1
            items.append(f"  [yellow]○[/yellow] {stripped[6:].strip()}")
        elif stripped.startswith("# ") and not items:
            # Show the title
            console.print(f"  [bold]{stripped[2:]}[/bold]")

    for item in items:
        console.print(item)

    total = done + pending
    if total:
        console.print(f"  [dim]{done}/{total} done[/dim]")


class ChatSession:
    """Interactive chat session over WebSocket.

    Each user query opens a new WebSocket connection to the /start endpoint.
    The server processes the query, streams back events, then closes the connection.
    The session loops, opening a new connection for each subsequent query.
    """

    def __init__(
        self,
        base_url: str,
        conversation_id: str,
        token: str,
        verify_ssl: bool = True,
        debug: bool = False,
        on_first_message: Callable[[str], Coroutine[Any, Any, None]] | None = None,
    ):
        self.base_url = base_url
        self.conversation_id = conversation_id
        self.token = token
        self.verify_ssl = verify_ssl
        self.debug = debug
        self.ssl_context = _build_ssl_context(verify_ssl)
        self.pending_file_ids: list[str] = []
        self._on_first_message = on_first_message
        self._message_count = 0

    async def run(self):
        """Run the interactive chat REPL (Read-Eval-Print Loop)."""
        console.print(
            Panel(
                f"[bold]Navari Chat[/bold] - {self.base_url}\n"
                f"Conversation: {self.conversation_id}\n\n"
                "Type your message and press Enter to send.\n"
                "Type [bold]/attach <path>[/bold] to attach a file.\n"
                "Type [bold]/files[/bold] to list attached files.\n"
                "Type [bold]/quit[/bold] or [bold]/q[/bold] to exit.\n"
                "Press [bold]Ctrl+C[/bold] to cancel a running query.",
                title="Chat Session",
                border_style="blue",
            )
        )

        while True:
            try:
                # Show attachment indicator in prompt
                if self.pending_file_ids:
                    prompt = f"\n[bold green]You ({len(self.pending_file_ids)} file(s))>[/bold green] "
                else:
                    prompt = "\n[bold green]You>[/bold green] "
                user_input = console.input(prompt).strip()
                if not user_input:
                    continue
                if user_input.lower() in ("/quit", "/q", "/exit"):
                    console.print(f"[dim]Ending chat session ({self.conversation_id}).[/dim]")
                    break

                # Handle slash commands
                if user_input.startswith("/attach "):
                    await self._handle_attach(user_input[8:].strip())
                    continue
                elif user_input.lower() == "/files":
                    self._handle_files_list()
                    continue
                elif user_input.lower() == "/clear":
                    self.pending_file_ids.clear()
                    console.print("[dim]Cleared attached files.[/dim]")
                    continue

                # Send query with any pending file attachments
                file_ids = list(self.pending_file_ids)
                self.pending_file_ids.clear()
                await self._send_query(user_input, file_ids=file_ids)

                # Auto-rename callback after first message
                self._message_count += 1
                if self._message_count == 1 and self._on_first_message:
                    try:
                        await self._on_first_message(user_input)
                    except Exception:
                        pass  # Non-critical

            except KeyboardInterrupt:
                console.print("\n[dim]Ending chat session.[/dim]")
                break
            except EOFError:
                console.print("\n[dim]Ending chat session.[/dim]")
                break

    async def _handle_attach(self, file_path_str: str):
        """Upload a file and stage it for the next message."""
        path = Path(file_path_str)
        if not path.exists():
            console.print(f"[red]Error:[/red] File not found: {file_path_str}")
            return

        try:
            async with NavariAPIClient(self.base_url, self.token, verify_ssl=self.verify_ssl) as client:
                result = await client.upload_file(self.conversation_id, path)
            file_id = result.get("id", result.get("file_id"))
            self.pending_file_ids.append(file_id)
            console.print(f"[green]✓[/green] Attached: [bold]{path.name}[/bold]  [dim]({file_id})[/dim]")
        except Exception as e:
            console.print(f"[red]Error:[/red] Upload failed: {e}")

    def _handle_files_list(self):
        """Show currently attached files."""
        if not self.pending_file_ids:
            console.print("[dim]No files attached. Use /attach <path> to add one.[/dim]")
        else:
            console.print(f"[bold]Attached files ({len(self.pending_file_ids)}):[/bold]")
            for fid in self.pending_file_ids:
                console.print(f"  {fid}")

    async def _send_query(self, query: str, file_ids: list[str] | None = None):
        """Send a single query over a new WebSocket connection and stream the response."""
        ws_url = _build_ws_url(self.base_url, self.conversation_id, self.token)

        if self.debug:
            console.print(f"[dim]Connecting to {ws_url[:80]}...[/dim]")

        # For wss://: use custom no-verify context, or True for default SSL (with truststore)
        # For ws://: no SSL needed
        ssl_ctx = (self.ssl_context or True) if ws_url.startswith("wss://") else None

        try:
            async with websockets.connect(
                ws_url,
                ssl=ssl_ctx,
                additional_headers={"Origin": self.base_url},
                open_timeout=60,
                close_timeout=10,
                max_size=10 * 1024 * 1024,  # 10MB max message size
            ) as ws:
                await self._handle_connection(ws, query, file_ids=file_ids)
        except websockets.exceptions.InvalidStatus as e:
            if e.response.status_code == 403:
                console.print("[red]Error:[/red] Access denied. Your session may have expired.")
                console.print("Try logging in again with [bold]navari login[/bold]")
            else:
                console.print(f"[red]Error:[/red] Connection failed (HTTP {e.response.status_code})")
            if self.debug:
                console.print(f"[dim]{e}[/dim]")
        except (websockets.exceptions.ConnectionClosed, ConnectionRefusedError) as e:
            console.print(f"[red]Error:[/red] Connection lost: {e}")
        except TimeoutError:
            console.print("[red]Error:[/red] Connection timed out")
        except Exception as e:
            console.print(f"[red]Error:[/red] {e}")
            if self.debug:
                traceback.print_exc()

    async def _handle_connection(self, ws, query: str, file_ids: list[str] | None = None):
        """Handle a single WebSocket connection lifecycle."""
        # Wait for initializing and connection_established events
        connected = await self._wait_for_connection(ws)
        if not connected:
            return

        # Send the query with optional file attachments
        msg = {"query": query}
        if file_ids:
            msg["file_ids"] = file_ids
        await ws.send(json.dumps(msg))

        # Stream the response
        await self._stream_response(ws)

    async def _wait_for_connection(self, ws) -> bool:
        """Wait for the server to finish initializing and send connection_established."""
        while True:
            try:
                raw = await asyncio.wait_for(ws.recv(), timeout=120)
                event = json.loads(raw)
                event_kind = event.get("event_kind", "")

                if self.debug:
                    console.print(f"[dim]<- {event_kind}[/dim]")

                if event_kind == "initializing":
                    console.print("[dim]Initializing agent...[/dim]", end="\r")
                    continue
                elif event_kind == "connection_established":
                    # Clear the initializing message
                    console.print(" " * 40, end="\r")
                    return True
                elif event_kind == "error":
                    console.print(f"[red]Error:[/red] {event.get('message', 'Unknown error')}")
                    return False
                else:
                    if self.debug:
                        console.print(f"[dim]Unexpected event during init: {event_kind}[/dim]")
                    continue

            except TimeoutError:
                console.print("[red]Error:[/red] Timed out waiting for agent initialization")
                return False

    async def _send_response_with_reconnect(self, ws, response: str):
        """Send user response, reconnecting via /attach if the WS dropped.

        During waiting_for_input the WebSocket may time out (proxy/infra).
        If sending fails, reconnect to the existing run via the /attach
        endpoint and resend.  Returns the (possibly new) ws.
        """
        payload = json.dumps({"response": response})
        try:
            await ws.send(payload)
            return ws
        except websockets.exceptions.ConnectionClosed:
            if self.debug:
                console.print("[dim]WS closed, reconnecting via /attach...[/dim]")
            else:
                console.print("[dim]Reconnecting...[/dim]")
            try:
                attach_url = _build_ws_attach_url(self.base_url, self.conversation_id, self.token)
                ssl_ctx = (self.ssl_context or True) if attach_url.startswith("wss://") else None
                new_ws = await websockets.connect(
                    attach_url,
                    ssl=ssl_ctx,
                    additional_headers={"Origin": self.base_url},
                    open_timeout=30,
                    close_timeout=10,
                    max_size=10 * 1024 * 1024,
                )
                # Reuse the same handshake logic as initial connections
                connected = await self._wait_for_connection(new_ws)
                if not connected:
                    console.print("[red]Error:[/red] Failed to reconnect")
                    return new_ws
                await new_ws.send(payload)
                return new_ws
            except Exception as e:
                console.print(f"[red]Error:[/red] Reconnection failed: {e}")
                raise

    async def _stream_response(self, ws):
        """Stream and display the agent's response events.

        The Navari agent communicates via tool calls:
        - send_message_to_user: agent's text output (message in args)
        - ask_user / ask_user_with_examples / ask_user_for_confirmation: user prompts
        - final_result: structured result at the end

        Tool call args are streamed via part_delta events with args_delta,
        then the complete call arrives as function_tool_call.
        Message content from send_message_to_user is streamed in real-time
        by extracting text from partial JSON as args_delta arrives.
        """
        # Track current tool call being streamed
        current_tool_name = ""
        tool_card_shown = False  # Whether (use tool "...") was printed for current call
        args_buffer = ""
        streamed_len = 0  # How many chars of message value already printed
        printed_header = False

        # send_message_to_user: content is streamed directly as text, no tool card
        HIDDEN_TOOLS = {"send_message_to_user"}
        # ask_user tools: show tool card + question + input prompt
        ASK_TOOLS = {"ask_user", "ask_user_with_examples", "ask_user_for_confirmation"}

        try:
            while True:
                try:
                    raw = await ws.recv()
                except websockets.exceptions.ConnectionClosed:
                    break  # Normal — server closes after query completes

                try:
                    event = json.loads(raw)
                except json.JSONDecodeError:
                    continue

                event_kind = event.get("event_kind", "")

                if self.debug and event_kind != "part_delta":
                    console.print(f"[dim]<- {event_kind}: {json.dumps(event, default=str)[:200]}[/dim]")

                # --- Custom events from server ---

                if event_kind == "processing_started":
                    printed_header = False
                    args_buffer = ""
                    current_tool_name = ""
                    streamed_len = 0

                elif event_kind == "processing_finished":
                    if printed_header:
                        console.print()  # Final newline after response

                elif event_kind == "processing_cancelled":
                    console.print("\n[yellow]Cancelled[/yellow]")

                elif event_kind == "agent_result":
                    summary = event.get("summary", "")
                    if summary:
                        console.print(f"\n[bold cyan]Navari>[/bold cyan] {summary}")

                elif event_kind == "error":
                    message = event.get("message", "Unknown error")
                    console.print(f"\n[red]Error:[/red] {message}")
                    break  # Server is done after an error

                elif event_kind == "waiting_for_input":
                    # Agent is asking for user input — prompt was already
                    # displayed from the function_tool_call event.
                    # Run input in an executor so the event loop stays alive
                    # for WebSocket ping/pong keepalive while the user types.
                    # If the WS still drops, _send_response_with_reconnect
                    # handles it by reconnecting via /attach.
                    try:
                        loop = asyncio.get_event_loop()
                        response = await loop.run_in_executor(
                            None, lambda: console.input("[bold yellow]>[/bold yellow] ").strip()
                        )
                        ws = await self._send_response_with_reconnect(ws, response)
                    except (KeyboardInterrupt, EOFError):
                        try:
                            await ws.send(json.dumps({"stop": True}))
                        except websockets.exceptions.ConnectionClosed:
                            pass
                        return

                elif event_kind == "tool_progress":
                    message = event.get("message", "")
                    if message:
                        console.print(f"  [dim]{message}[/dim]")

                elif event_kind == "compaction":
                    if self.debug:
                        orig = event.get("original_messages", "?")
                        compacted = event.get("compacted_messages", "?")
                        console.print(f"  [dim]\\[compaction: {orig} -> {compacted} messages][/dim]")

                elif event_kind == "mode_changed":
                    mode = event.get("mode", "unknown")
                    console.print(f"  [dim]\\[mode: {mode}][/dim]")

                # --- Pydantic-AI streaming events ---

                elif event_kind == "part_start":
                    # New tool call starting — extract tool name for streaming
                    part = event.get("part", {})
                    current_tool_name = part.get("tool_name", "")
                    args_buffer = ""
                    streamed_len = 0

                    # Show tool card for tools that aren't handled specially
                    if current_tool_name and current_tool_name not in HIDDEN_TOOLS:
                        console.print(f'  [dim](use tool "{current_tool_name}")[/dim]')
                        tool_card_shown = True

                elif event_kind == "part_delta":
                    delta = event.get("delta", {})
                    if isinstance(delta, dict):
                        # Tool call args streaming
                        args_delta = delta.get("args_delta", "")
                        if args_delta:
                            if isinstance(args_delta, str):
                                args_buffer += args_delta
                            else:
                                # dict delta — can't stream incrementally
                                args_buffer = json.dumps(args_delta)

                            # Stream message text in real-time for send_message_to_user
                            if current_tool_name == "send_message_to_user":
                                text = _extract_json_string_value(args_buffer, "message")
                                if text and len(text) > streamed_len:
                                    new_text = text[streamed_len:]
                                    if not printed_header:
                                        console.print("\n[bold cyan]Navari>[/bold cyan] ", end="")
                                        printed_header = True
                                    print(new_text, end="", flush=True)
                                    streamed_len = len(text)

                        # Direct text streaming (non-tool responses)
                        content = delta.get("content_delta", "")
                        if content:
                            if not printed_header:
                                console.print("\n[bold cyan]Navari>[/bold cyan] ", end="")
                                printed_header = True
                            print(content, end="", flush=True)

                elif event_kind == "part_end":
                    # Finish streaming — add newline if we streamed message content
                    if current_tool_name == "send_message_to_user" and streamed_len > 0:
                        print(flush=True)  # Final newline after streamed message
                    args_buffer = ""
                    current_tool_name = ""
                    streamed_len = 0

                elif event_kind == "function_tool_call":
                    part = event.get("part", {})
                    tool_name = part.get("tool_name", "")
                    args_str = part.get("args", "")

                    # Parse tool arguments (may be a JSON string or already a dict)
                    if isinstance(args_str, dict):
                        args = args_str
                    elif args_str:
                        try:
                            args = json.loads(args_str)
                        except (json.JSONDecodeError, TypeError):
                            args = {}
                    else:
                        args = {}

                    if tool_name == "send_message_to_user":
                        # If already streamed via part_delta, skip.
                        # Otherwise display the full message (e.g. test models
                        # that send complete events without streaming).
                        if streamed_len == 0:
                            message = args.get("message", "")
                            if message:
                                if not printed_header:
                                    console.print("\n[bold cyan]Navari>[/bold cyan] ", end="")
                                    printed_header = True
                                console.print(Markdown(message))

                    elif tool_name in ASK_TOOLS:
                        # Display the question the agent is asking
                        question = args.get("question", "")
                        if question:
                            if not printed_header:
                                printed_header = True
                            console.print(f"\n[bold cyan]Navari>[/bold cyan] {question}")
                        # Show example responses if provided (ask_user_with_examples)
                        examples = args.get("examples", [])
                        for i, example in enumerate(examples, 1):
                            console.print(f"  [dim]{i}. {example}[/dim]")

                    elif tool_name == "write_to_do_file":
                        # Show compact todo list summary
                        content = args.get("content", "")
                        if content:
                            _print_todo_summary(content)

                    elif tool_name == "report_file_created":
                        # Show filename from args
                        fp = args.get("file_path", "")
                        name = fp.rsplit("/", 1)[-1] if fp else ""
                        console.print(f"  [dim](file created: {name or fp})[/dim]")

                    elif not tool_card_shown:
                        # Only show if not already displayed from part_start
                        console.print(f'  [dim](use tool "{tool_name}")[/dim]')

                    tool_card_shown = False

                elif event_kind == "function_tool_result":
                    pass

                elif event_kind == "final_result":
                    pass

                elif self.debug:
                    console.print(f"  [dim]\\[unhandled: {event_kind}][/dim]")

        except KeyboardInterrupt:
            console.print("\n[yellow]Query cancelled[/yellow]")
            try:
                await ws.send(json.dumps({"stop": True}))
            except Exception:
                pass
