"""SSE chat client for interactive conversations with Navari agent via AG-UI protocol."""

import asyncio
import json
import tempfile
import traceback
from collections.abc import Callable, Coroutine
from pathlib import Path
from typing import Any
from uuid import uuid4

import httpx
from ag_ui.core import EventType, UserMessage
from rich.console import Console
from rich.panel import Panel

from navari_cli.api_client import NavariAPIClient
from navari_cli.file_utils import (
    ASK_TOOLS,
    MESSAGE_TOOLS,
    find_workspace_file,
    open_url_in_browser,
    open_with_system_app,
    unique_path,
)

console = Console()


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:
            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 AG-UI SSE protocol.

    Each user query sends an HTTP POST to the AG-UI endpoint and streams the
    SSE response.  When the agent needs user input (ask_user), the run finishes
    with DeferredToolRequests; the CLI prompts the user and sends a new POST
    with the approval response.
    """

    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.pending_file_ids: list[str] = []
        self.workspace_files: list[dict[str, str]] = []  # [{name, path, url}] from report_file_created
        self._on_first_message = on_first_message
        self._message_count = 0
        self._messages: list[dict] = []  # AG-UI message history for the thread

    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]/browse <name>[/bold] to open a workspace file in the browser.\n"
                "Type [bold]/open <name>[/bold] to open with the system default app.\n"
                "Type [bold]/download <name>[/bold] to download a workspace file.\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:
                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
                elif user_input.startswith("/browse "):
                    await self._handle_browse(user_input[8:].strip())
                    continue
                elif user_input.startswith("/open "):
                    await self._handle_open(user_input[6:].strip())
                    continue
                elif user_input.startswith("/download "):
                    await self._handle_download(user_input[10:].strip())
                    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

    # --- File handling (attach, browse, open, download) ---

    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}")

    def _resolve_file_url(self, name_or_url: str) -> str | None:
        """Resolve a filename, path, or URL to a full download URL."""
        if name_or_url.startswith(("http://", "https://")):
            return name_or_url
        return find_workspace_file(self.workspace_files, name_or_url)

    def _print_available_files(self):
        """Print list of available workspace files."""
        if self.workspace_files:
            console.print("[dim]Available files:[/dim]")
            for wf in self.workspace_files:
                console.print(f"  [dim]{wf['path'] or wf['name']}[/dim]")

    async def _fetch_file(self, name_or_url: str) -> tuple[str, bytes] | None:
        """Resolve and download a workspace file. Returns (filename, content) or None."""
        url = self._resolve_file_url(name_or_url)
        if not url:
            console.print(f"[red]Error:[/red] File not found: {name_or_url}")
            self._print_available_files()
            return None
        filename = url.rsplit("/", 1)[-1] or "download"
        async with httpx.AsyncClient(verify=self.verify_ssl) as http:
            resp = await http.get(url, headers={"Authorization": f"Bearer {self.token}"})
            resp.raise_for_status()
        return filename, resp.content

    async def _handle_browse(self, name_or_url: str):
        """Open a workspace file in the browser."""
        url = self._resolve_file_url(name_or_url)
        if not url:
            console.print(f"[red]Error:[/red] File not found: {name_or_url}")
            self._print_available_files()
            return
        open_url_in_browser(url, self.token)
        console.print("[green]✓[/green] Opened in browser")

    async def _handle_open(self, name_or_url: str):
        """Download and open a workspace file with the system default app."""
        try:
            result = await self._fetch_file(name_or_url)
            if not result:
                return
            filename, content = result
            tmp_dir = Path(tempfile.mkdtemp())
            out_path = tmp_dir / filename
            out_path.write_bytes(content)
            open_with_system_app(out_path)
            console.print(f"[green]✓[/green] Opened: [bold]{filename}[/bold]")
        except Exception as e:
            console.print(f"[red]Error:[/red] Failed to open: {e}")

    async def _handle_download(self, name_or_url: str):
        """Download a workspace file to the current directory."""
        try:
            result = await self._fetch_file(name_or_url)
            if not result:
                return
            filename, content = result
            out_path = unique_path(Path(filename))
            out_path.write_bytes(content)
            console.print(f"[green]✓[/green] Downloaded: [bold]{out_path}[/bold] ({len(content):,} bytes)")
        except Exception as e:
            console.print(f"[red]Error:[/red] Download failed: {e}")

    # --- AG-UI SSE communication ---

    def _agui_url(self) -> str:
        """Build the AG-UI endpoint URL."""
        return f"{self.base_url}/api/conversations/{self.conversation_id}/agui"

    def _auth_headers(self) -> dict[str, str]:
        """Build authorization headers."""
        return {"Authorization": f"Bearer {self.token}"}

    def _build_run_input(self, messages: list, forwarded_props: dict | None = None) -> dict:
        """Build RunAgentInput body as a dict for the AG-UI POST request.

        Uses camelCase keys to match the canonical AG-UI JSON format.
        """
        return {
            "threadId": self.conversation_id,
            "runId": str(uuid4()),
            "messages": [m.model_dump() if hasattr(m, "model_dump") else m for m in messages],
            "state": {},
            "tools": [],
            "context": [],
            "forwardedProps": forwarded_props or {},
        }

    async def _send_query(self, query: str, file_ids: list[str] | None = None):
        """Send a query via AG-UI protocol and stream the SSE response."""
        user_msg = UserMessage(id=str(uuid4()), role="user", content=query)
        self._messages.append(user_msg.model_dump())

        body = self._build_run_input(
            messages=[user_msg],
            forwarded_props={"file_ids": file_ids} if file_ids else None,
        )

        if self.debug:
            console.print(f"[dim]POST {self._agui_url()}[/dim]")

        try:
            async with httpx.AsyncClient(verify=self.verify_ssl, timeout=httpx.Timeout(300.0)) as client:
                async with client.stream(
                    "POST",
                    self._agui_url(),
                    json=body,
                    headers=self._auth_headers(),
                ) as response:
                    if 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]")
                        return
                    if response.status_code == 409:
                        console.print("[red]Error:[/red] Another request is already processing.")
                        return
                    if response.status_code != 200:
                        console.print(f"[red]Error:[/red] Server returned HTTP {response.status_code}")
                        return

                    console.print("[dim]Initializing agent...[/dim]", end="\r")
                    await self._stream_sse_response(response)

        except httpx.ConnectError:
            console.print("[red]Error:[/red] Could not connect to server")
        except httpx.TimeoutException:
            console.print("[red]Error:[/red] Request timed out")
        except KeyboardInterrupt:
            console.print("\n[yellow]Query cancelled[/yellow]")
            # Cancel the active run
            try:
                async with httpx.AsyncClient(verify=self.verify_ssl) as client:
                    await client.post(
                        f"{self.base_url}/api/runs/{self.conversation_id}/cancel",
                        headers=self._auth_headers(),
                    )
            except Exception:
                pass
        except Exception as e:
            console.print(f"[red]Error:[/red] {e}")
            if self.debug:
                traceback.print_exc()

    async def _send_approval(self, tool_call_id: str, response: str, file_ids: list[str] | None = None):
        """Send an approval response for a deferred tool request.

        Waits briefly before sending to allow the server to finish cleaning up
        the previous run's MCP sessions, then retries with exponential backoff
        on 500/409 errors.
        """
        # Give the server time to clean up MCP sessions from the previous run
        await asyncio.sleep(1)

        user_msg = UserMessage(id=str(uuid4()), role="user", content=response)
        self._messages.append(user_msg.model_dump())

        forwarded_props = {"approval_response_for": tool_call_id}
        if file_ids:
            forwarded_props["file_ids"] = file_ids

        body = self._build_run_input(messages=[user_msg], forwarded_props=forwarded_props)

        if self.debug:
            console.print(f"[dim]POST {self._agui_url()} (approval for {tool_call_id})[/dim]")

        max_retries = 4
        for attempt in range(max_retries):
            async with httpx.AsyncClient(verify=self.verify_ssl, timeout=httpx.Timeout(300.0)) as client:
                async with client.stream(
                    "POST",
                    self._agui_url(),
                    json=body,
                    headers=self._auth_headers(),
                ) as resp:
                    if resp.status_code in (500, 409) and attempt < max_retries - 1:
                        delay = 2 ** (attempt + 1)  # 2, 4, 8 seconds
                        if self.debug:
                            console.print(f"[dim]Server {resp.status_code}, retrying in {delay}s...[/dim]")
                        await asyncio.sleep(delay)
                        continue
                    if resp.status_code != 200:
                        console.print(f"[red]Error:[/red] Approval failed (HTTP {resp.status_code})")
                        return
                    await self._stream_sse_response(resp)
                    return

    async def _stream_sse_response(self, response):
        """Stream and display AG-UI SSE events from an httpx response.

        AG-UI events replace the old pydantic-ai WebSocket events:
        - TEXT_MESSAGE_START/CONTENT/END → agent text (replaces send_message_to_user)
        - TOOL_CALL_START/ARGS/END/RESULT → tool execution
        - RUN_FINISHED with DeferredToolRequests → ask_user approval flow
        - RUN_ERROR → error display
        - CUSTOM → mode changes, compaction
        """
        printed_header = False
        current_tool_name = ""
        args_buffer = ""
        pending_approval = None  # (tool_call_id, tool_name, metadata) if approval needed

        async for line in response.aiter_lines():
            if not line.startswith("data: "):
                continue  # Skip keepalive comments and empty lines

            try:
                event = json.loads(line[6:])
            except json.JSONDecodeError:
                continue

            event_type = event.get("type", "")

            if self.debug and event_type not in ("TEXT_MESSAGE_CONTENT", "TOOL_CALL_ARGS"):
                console.print(f"[dim]<- {event_type}: {json.dumps(event, default=str)[:200]}[/dim]")

            # --- Run lifecycle ---

            if event_type == EventType.RUN_STARTED.value:
                printed_header = False
                console.print(" " * 40, end="\r")  # Clear "Initializing agent..."

            elif event_type == EventType.RUN_FINISHED.value:
                if printed_header:
                    console.print()  # Final newline

                # Check result for summary or approvals
                result = event.get("result")
                if isinstance(result, dict):
                    # Show final result summary
                    summary = result.get("summary")
                    if summary:
                        console.print(f"\n[bold cyan]Navari>[/bold cyan] {summary}")

                    # Check for approvals (ask_user flow)
                    deferred = result.get("approvals")
                    if deferred:
                        for req in deferred:
                            tool_call_id = req.get("tool_call_id", "")
                            tool_name = req.get("tool_name", "")
                            metadata = req.get("metadata", {})
                            pending_approval = (tool_call_id, tool_name, metadata)
                            break  # Handle first approval

                break  # Run complete, stop reading SSE stream

            elif event_type == EventType.RUN_ERROR.value:
                message = event.get("message", "Unknown error")
                console.print(f"\n[red]Error:[/red] {message}")
                break  # Run failed, stop reading SSE stream

            # --- Text messages (replaces send_message_to_user) ---

            elif event_type == EventType.TEXT_MESSAGE_START.value:
                if not printed_header:
                    console.print("\n[bold cyan]Navari>[/bold cyan] ", end="")
                    printed_header = True

            elif event_type == EventType.TEXT_MESSAGE_CONTENT.value:
                delta = event.get("delta", "")
                if delta:
                    if not printed_header:
                        console.print("\n[bold cyan]Navari>[/bold cyan] ", end="")
                        printed_header = True
                    print(delta, end="", flush=True)

            elif event_type == EventType.TEXT_MESSAGE_END.value:
                if printed_header:
                    print(flush=True)  # Final newline
                    printed_header = False

            # --- Tool calls ---

            elif event_type == EventType.TOOL_CALL_START.value:
                current_tool_name = event.get("toolCallName", "")
                args_buffer = ""
                # send_message_to_user: message shown as Navari> text, no tool card
                if current_tool_name and current_tool_name not in MESSAGE_TOOLS:
                    console.print(f'  [dim](use tool "{current_tool_name}")[/dim]')

            elif event_type == EventType.TOOL_CALL_ARGS.value:
                delta = event.get("delta", "")
                if delta:
                    args_buffer += delta

            elif event_type == EventType.TOOL_CALL_END.value:
                # Process buffered args for special tools
                if current_tool_name in MESSAGE_TOOLS and args_buffer:
                    try:
                        args = json.loads(args_buffer)
                        message = args.get("message", "")
                        if message:
                            if not printed_header:
                                console.print("\n[bold cyan]Navari>[/bold cyan] ", end="")
                                printed_header = True
                            print(message, flush=True)
                    except json.JSONDecodeError:
                        pass

                elif current_tool_name in ASK_TOOLS and args_buffer:
                    try:
                        args = json.loads(args_buffer)
                        question = args.get("question", "")
                        if question:
                            console.print(f"\n[bold cyan]Navari>[/bold cyan] {question}")
                        examples = args.get("examples", [])
                        for i, example in enumerate(examples, 1):
                            console.print(f"  [dim]{i}. {example}[/dim]")
                    except json.JSONDecodeError:
                        pass

                elif current_tool_name == "write_to_do_file" and args_buffer:
                    try:
                        args = json.loads(args_buffer)
                        content = args.get("content", "")
                        if content:
                            _print_todo_summary(content)
                    except json.JSONDecodeError:
                        pass

                elif current_tool_name == "report_file_created" and args_buffer:
                    try:
                        args = json.loads(args_buffer)
                        fp = args.get("file_path", "")
                        name = fp.rsplit("/", 1)[-1] if fp else ""
                        console.print(f"  [dim](file created: {name or fp})[/dim]")
                    except json.JSONDecodeError:
                        pass

                current_tool_name = ""
                args_buffer = ""

            elif event_type == EventType.TOOL_CALL_RESULT.value:
                result_str = event.get("content", event.get("result", ""))
                try:
                    result_content = json.loads(result_str) if isinstance(result_str, str) else result_str
                except (json.JSONDecodeError, TypeError):
                    result_content = {}
                if isinstance(result_content, dict):
                    # Track workspace files from report_file_created results
                    if result_content.get("file_url"):
                        name = result_content.get("file_name", "")
                        path = result_content.get("display_path", "")
                        url = result_content["file_url"]
                        if url.startswith("/"):
                            url = f"{self.base_url.rstrip('/')}{url}"
                        self.workspace_files.append({"name": name, "path": path, "url": url})
                        if path:
                            console.print(f"    [dim]{path}[/dim]")

            # --- Custom events ---

            elif event_type == EventType.CUSTOM.value:
                custom_name = event.get("name", "")
                custom_value = event.get("value", {})

                if custom_name == "mode_changed":
                    mode = custom_value.get("mode", "unknown") if isinstance(custom_value, dict) else custom_value
                    console.print(f"  [dim]\\[mode: {mode}][/dim]")
                elif custom_name == "compaction" and self.debug:
                    if isinstance(custom_value, dict):
                        orig = custom_value.get("original_messages", "?")
                        compacted = custom_value.get("compacted_messages", "?")
                        console.print(f"  [dim]\\[compaction: {orig} -> {compacted} messages][/dim]")
                elif self.debug:
                    console.print(f"  [dim]\\[custom: {custom_name}][/dim]")

            elif event_type == EventType.ACTIVITY_SNAPSHOT.value:
                # A2A agent progress
                activity = event.get("activity", {})
                if isinstance(activity, dict):
                    title = activity.get("title", "")
                    if title:
                        console.print(f"  [dim]{title}[/dim]")

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

        # Close the SSE response before blocking on user input (approval flow)
        await response.aclose()

        # Handle pending approval (ask_user flow)
        if pending_approval:
            tool_call_id, tool_name, metadata = pending_approval
            try:
                response_text = console.input("[bold yellow]>[/bold yellow] ").strip()
                await self._send_approval(tool_call_id, response_text)
            except (KeyboardInterrupt, EOFError):
                return
