import asyncio
import uuid
from typing import Any

import httpx
from rich.console import Console
from rich.markdown import Markdown
from rich.prompt import Prompt

try:
    from fasta2a.client import A2AClient as FastA2AClient
    from fasta2a.schema import Message, TextPart
except ImportError as e:
    raise ImportError(
        "A2A support requires the 'a2a' extra. "
        "Install with: uv pip install 'haiku.rag[a2a]'"
    ) from e


class A2AClient:
    """Interactive A2A protocol client."""

    def __init__(self, base_url: str = "http://localhost:8000"):
        """Initialize A2A client.

        Args:
            base_url: Base URL of the A2A server
        """
        self.base_url = base_url.rstrip("/")
        http_client = httpx.AsyncClient(timeout=60.0)
        self._client = FastA2AClient(base_url=base_url, http_client=http_client)

    async def close(self):
        """Close the HTTP client."""
        await self._client.http_client.aclose()

    async def get_agent_card(self) -> dict[str, Any]:
        """Fetch the agent card from the A2A server.

        Returns:
            Agent card dictionary with agent capabilities and metadata
        """
        response = await self._client.http_client.get(
            f"{self.base_url}/.well-known/agent-card.json"
        )
        response.raise_for_status()
        return response.json()

    async def send_message(
        self,
        text: str,
        context_id: str | None = None,
        skill_id: str | None = None,
    ) -> dict[str, Any]:
        """Send a message to the A2A agent and wait for completion.

        Args:
            text: Message text to send
            context_id: Optional conversation context ID (creates new if None)
            skill_id: Optional skill ID to use (defaults to document-qa)

        Returns:
            Completed task with response messages and artifacts
        """
        if context_id is None:
            context_id = str(uuid.uuid4())

        message = Message(
            kind="message",
            role="user",
            message_id=str(uuid.uuid4()),
            parts=[TextPart(kind="text", text=text)],
        )

        metadata: dict[str, Any] = {"contextId": context_id}
        if skill_id:
            metadata["skillId"] = skill_id

        response = await self._client.send_message(message, metadata=metadata)

        if "error" in response:
            return {"error": response["error"]}

        result = response.get("result")
        if not result:
            return {"result": result}

        # Result can be either Task or Message - check if it's a Task with an id
        if result.get("kind") == "task":
            task_id = result.get("id")
            if task_id:
                # Poll for task completion
                return await self.wait_for_task(task_id)

        # Return the message directly
        return {"result": result}

    async def wait_for_task(
        self, task_id: str, max_wait: int = 120, poll_interval: float = 0.5
    ) -> dict[str, Any]:
        """Poll for task completion.

        Args:
            task_id: Task ID to poll for
            max_wait: Maximum time to wait in seconds
            poll_interval: Interval between polls in seconds

        Returns:
            Completed task result
        """
        import time

        start_time = time.time()

        while time.time() - start_time < max_wait:
            task_response = await self._client.get_task(task_id)

            if "error" in task_response:
                return {"error": task_response["error"]}

            task = task_response.get("result")
            if not task:
                raise Exception("No task in response")

            state = task.get("status", {}).get("state")

            if state == "completed":
                return {"result": task}
            elif state == "failed":
                raise Exception(f"Task failed: {task}")

            await asyncio.sleep(poll_interval)

        raise TimeoutError(f"Task {task_id} did not complete within {max_wait}s")


def print_agent_card(card: dict[str, Any], console: Console):
    """Pretty print the agent card using Rich."""
    console.print()
    console.print("[bold]Agent Card[/bold]")
    console.rule()

    console.print(f"  [repr.attrib_name]name[/repr.attrib_name]: {card.get('name')}")
    console.print(
        f"  [repr.attrib_name]description[/repr.attrib_name]: {card.get('description')}"
    )
    console.print(
        f"  [repr.attrib_name]version[/repr.attrib_name]: {card.get('version')}"
    )
    console.print(
        f"  [repr.attrib_name]protocol version[/repr.attrib_name]: {card.get('protocolVersion')}"
    )

    skills = card.get("skills", [])
    console.print(f"\n[bold cyan]Skills ({len(skills)}):[/bold cyan]")
    for skill in skills:
        console.print(f"  • {skill.get('id')}: {skill.get('name')}")
        console.print(f"    [dim]{skill.get('description')}[/dim]")
        examples = skill.get("examples", [])
        if examples:
            console.print(f"    [dim]Examples: {', '.join(examples[:2])}[/dim]")
    console.print()


def print_response(response: dict[str, Any], console: Console):
    """Pretty print the A2A response using Rich."""
    if "error" in response:
        console.print(f"[red]Error: {response['error']}[/red]")
        return

    result = response.get("result", {})

    # Get messages from history and artifacts from completed task
    history = result.get("history", [])
    artifacts = result.get("artifacts", [])

    # Print agent messages from history with markdown rendering
    for msg in history:
        if msg.get("role") == "agent":
            for part in msg.get("parts", []):
                if part.get("kind") == "text":
                    text = part.get("text", "")
                    # Render as markdown
                    console.print()
                    console.print("[bold green]Answer:[/bold green]")
                    console.print(Markdown(text))

    # Print artifacts summary with details
    if artifacts:
        console.rule("[dim]Artifacts generated[/dim]")
        summary_lines = []

        for artifact in artifacts:
            name = artifact.get("name", "")
            parts = artifact.get("parts", [])

            if name == "search_results" and parts:
                data = parts[0].get("data", {})
                query = data.get("query", "")
                results = data.get("results", [])
                summary_lines.append(f"🔍 search: '{query}' ({len(results)} results)")

            elif name == "document" and parts:
                part = parts[0]
                if part.get("kind") == "text":
                    text = part.get("text", "")
                    length = len(text)
                    summary_lines.append(f"📄 document ({length} chars)")

            elif name == "qa_result" and parts:
                data = parts[0].get("data", {})
                skill = data.get("skill", "unknown")
                summary_lines.append(f"💬 {skill}")

        if summary_lines:
            console.print(f"[dim]{' • '.join(summary_lines)}[/dim]")

    console.print()


async def run_interactive_client(url: str = "http://localhost:8000"):
    """Run the interactive A2A client.

    Args:
        url: Base URL of the A2A server
    """
    console = Console()
    client = A2AClient(url)

    console.print("[bold]haiku.rag A2A interactive client[/bold]")
    console.print()

    # Fetch and display agent card
    console.print("[dim]Fetching agent card...[/dim]")
    try:
        card = await client.get_agent_card()
        print_agent_card(card, console)
    except Exception as e:
        console.print(f"[red]Error fetching agent card: {e}[/red]")
        await client.close()
        return

    # Create a conversation context
    context_id = str(uuid.uuid4())
    console.print(f"[dim]context id: {context_id}[/dim]")
    console.print("[dim]Type your questions (or 'quit' to exit)[/dim]\n")

    try:
        while True:
            try:
                question = Prompt.ask("[bold blue]Question[/bold blue]").strip()
                if not question:
                    continue

                if question.lower() in ("quit", "exit", "q"):
                    console.print("\n[dim]Goodbye![/dim]")
                    break

                response = await client.send_message(question, context_id=context_id)
                print_response(response, console)

            except KeyboardInterrupt:
                console.print("\n\n[dim]Exiting...[/dim]")
                break
            except Exception as e:
                console.print(f"\n[red]Error: {e}[/red]\n")
    finally:
        await client.close()
