"""A2A Protocol Transport - FastA2A integration.

This module provides A2A (Agent-to-Agent) protocol support for Albus,
enabling interoperability with other A2A-compliant agent frameworks
(PydanticAI, LangGraph, etc.).

Usage:
    # With optional a2a dependencies installed:
    # uv pip install albusos[a2a]

    from server.transport.a2a import create_a2a_app, AlbusWorker

    # Create A2A ASGI app
    app = create_a2a_app(runtime)

    # Run with uvicorn alongside main server
    # Or mount as sub-application

The A2A protocol is a standard for agent-to-agent communication over HTTP.
See: https://a2aproject.github.io/A2A/

Architecture:
    External A2A Client
           │
           ▼
    FastA2A ASGI App (JSON-RPC over HTTP)
           │
           ▼
    AlbusWorker (bridges A2A Tasks to Albus)
           │
           ▼
    Runtime.turn()
           │
           ▼
    PathwayService → PathwayVM
"""

from __future__ import annotations

import json
import logging
import uuid
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Literal

from pydantic import BaseModel, Field

from core.common import JsonObject
from core.types.messages import ConversationMessage

logger = logging.getLogger(__name__)

# Lazy imports for optional dependencies
_FASTA2A_AVAILABLE = False
try:
    from fasta2a import FastA2A, Worker
    from fasta2a.broker import Broker, InMemoryBroker
    from fasta2a.schema import (
        DataPart,
        Part,
        TaskIdParams,
        TaskSendParams,
        TextPart,
    )
    from fasta2a.storage import InMemoryStorage, Storage

    _FASTA2A_AVAILABLE = True
except ImportError:
    pass

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from server.runtime import ExecutionResult, Runtime


# =============================================================================
# A2A PYDANTIC MODELS — typed replacements for fasta2a TypedDicts
# =============================================================================


class A2ATextPart(BaseModel):
    """A text part within an A2A message or artifact."""

    kind: Literal["text"] = "text"
    text: str


class A2ADataPart(BaseModel):
    """A data part within an A2A message or artifact."""

    kind: Literal["data"] = "data"
    data: JsonObject = Field(default_factory=dict)


class A2AMessage(BaseModel):
    """An A2A protocol message (replaces fasta2a Message TypedDict)."""

    role: Literal["user", "agent"]
    parts: list[A2ATextPart | A2ADataPart]
    kind: Literal["message"] = "message"
    message_id: str = Field(default_factory=lambda: f"msg_{uuid.uuid4().hex[:12]}")


class A2AArtifact(BaseModel):
    """An A2A protocol artifact (replaces fasta2a Artifact TypedDict)."""

    artifact_id: str = Field(default_factory=lambda: f"art_{uuid.uuid4().hex[:12]}")
    name: str
    description: str = ""
    parts: list[A2ATextPart | A2ADataPart] = Field(default_factory=list)
    metadata: JsonObject = Field(default_factory=dict)


class A2ASkill(BaseModel):
    """An A2A protocol skill (replaces fasta2a Skill TypedDict)."""

    id: str
    name: str
    description: str = ""
    tags: list[str] = Field(default_factory=list)
    input_modes: list[str] = Field(default_factory=list)
    output_modes: list[str] = Field(default_factory=list)


def is_a2a_available() -> bool:
    """Check if A2A dependencies are installed."""
    return _FASTA2A_AVAILABLE


# =============================================================================
# ALBUS WORKER
# =============================================================================


@dataclass
class AlbusWorker:
    """A2A Worker that bridges to Runtime.

    This worker implements the FastA2A Worker interface, translating
    A2A protocol Tasks into Albus agent turns.

    Mapping:
        A2A Task → Albus Thread
        A2A context_id → Albus thread_id
        A2A Message → Albus message
        A2A Artifact → Albus response

    Args:
        runtime: The Runtime instance
        broker: FastA2A broker for task scheduling
        storage: FastA2A storage for task/context persistence
        default_agent_id: Default agent to use (default: "host")
    """

    runtime: Runtime
    broker: Any  # Broker type from fasta2a
    storage: Any  # Storage type from fasta2a
    default_agent_id: str = "host"

    # Internal state
    _running_tasks: dict[str, bool] = field(default_factory=dict)

    def __post_init__(self):
        if not _FASTA2A_AVAILABLE:
            raise ImportError(
                "FastA2A is not installed. Install with: uv pip install albusos[a2a]"
            )

    @asynccontextmanager
    async def run(self) -> AsyncIterator[None]:
        """Run the worker - receives and processes tasks from the broker."""
        import anyio

        async with anyio.create_task_group() as tg:
            tg.start_soon(self._task_loop)
            yield
            tg.cancel_scope.cancel()

    async def _task_loop(self) -> None:
        """Main loop - receives task operations from broker."""
        async for task_op in self.broker.receive_task_operations():
            await self._handle_task_operation(task_op)

    async def _handle_task_operation(self, task_op: dict[str, Any]) -> None:
        """Handle a task operation from the broker."""
        try:
            operation = task_op.get("operation")
            params = task_op.get("params", {})

            if operation == "run":
                await self.run_task(params)
            elif operation == "cancel":
                await self.cancel_task(params)
            else:
                logger.warning(f"Unknown task operation: {operation}")
        except Exception as e:
            logger.error(f"Task operation failed: {e}", exc_info=True)
            task_id = task_op.get("params", {}).get("id")
            if task_id:
                await self.storage.update_task(task_id, state="failed")

    async def run_task(self, params: dict[str, Any]) -> None:
        """Execute an A2A task by bridging to Runtime.

        This is the core bridge: it takes an A2A Task and runs it
        through the Albus execution engine.
        """
        task_id = params.get("id")
        if not task_id:
            logger.error("Task missing id")
            return

        # Load task from storage
        task = await self.storage.load_task(task_id)
        if not task:
            logger.error(f"Task not found: {task_id}")
            return

        # Mark as working
        await self.storage.update_task(task_id, state="working")
        self._running_tasks[task_id] = True

        try:
            # Extract context
            context_id = task.get("context_id", f"a2a-{uuid.uuid4().hex[:8]}")

            # Load conversation context (previous messages in this context)
            context = await self.storage.load_context(context_id) or []

            # Get the current message from task history — validate as models
            raw_history = task.get("history", [])
            if not raw_history:
                message_data = params.get("message", {})
                if message_data:
                    raw_history = [message_data]

            history = [A2AMessage.model_validate(m) for m in raw_history]

            # Extract user message text
            user_message = self._extract_message_text(history)
            if not user_message:
                err_msg = self._create_error_message("No message content provided")
                await self.storage.update_task(
                    task_id,
                    state="failed",
                    new_messages=[err_msg.model_dump(mode="json")],
                )
                return

            # Single agent — always route to host
            agent_id = self.default_agent_id

            logger.info(
                f"[A2A] Running task {task_id} with agent {agent_id}: {user_message[:100]}..."
            )

            # Execute via Albus
            # Map A2A context_id to Albus thread_id
            result = await self.runtime.turn(
                agent_id=agent_id,
                message=user_message,
                thread_id=context_id,  # Use A2A context as thread
            )

            # Build response message
            response_text = result.response
            success = result.success

            response_message = A2AMessage(
                role="agent",
                parts=[A2ATextPart(text=response_text)],
            )

            # Update context with new messages (dicts for fasta2a storage)
            context.extend(m.model_dump(mode="json") for m in history)
            context.append(response_message.model_dump(mode="json"))
            await self.storage.update_context(context_id, context)

            # Build artifacts from result
            artifacts = self.build_artifacts(result)

            # Update task as completed (convert models to dicts for fasta2a)
            await self.storage.update_task(
                task_id,
                state="completed" if success else "failed",
                new_messages=[response_message.model_dump(mode="json")],
                new_artifacts=[a.model_dump(mode="json") for a in artifacts],
            )

            logger.info(f"[A2A] Task {task_id} completed: success={success}")

        except Exception as e:
            logger.error(f"[A2A] Task {task_id} failed: {e}", exc_info=True)
            err_msg = self._create_error_message(str(e))
            await self.storage.update_task(
                task_id,
                state="failed",
                new_messages=[err_msg.model_dump(mode="json")],
            )
        finally:
            self._running_tasks.pop(task_id, None)

    async def cancel_task(self, params: dict[str, Any]) -> None:
        """Cancel a running task."""
        task_id = params.get("id")
        if not task_id:
            return

        # Mark task as cancelled in our tracking
        if task_id in self._running_tasks:
            self._running_tasks.pop(task_id, None)

        # Update storage
        await self.storage.update_task(task_id, state="canceled")
        logger.info(f"[A2A] Task {task_id} cancelled")

    def build_message_history(self, history: list[A2AMessage]) -> list[ConversationMessage]:
        """Convert A2A message history to Albus format.

        A2A Messages have parts (TextPart, DataPart).
        Albus expects ConversationMessage models.
        """
        messages: list[ConversationMessage] = []
        for msg in history:
            text_parts: list[str] = []
            for part in msg.parts:
                if isinstance(part, A2ATextPart):
                    text_parts.append(part.text)
                elif isinstance(part, A2ADataPart):
                    text_parts.append(json.dumps(part.data))

            content = "\n".join(text_parts)
            if content:
                messages.append(
                    ConversationMessage(
                        role="user" if msg.role == "user" else "assistant",
                        content=content,
                    )
                )

        return messages

    def build_artifacts(self, result: ExecutionResult) -> list[A2AArtifact]:
        """Convert Albus ExecutionResult to A2A Artifacts.

        Artifacts are the final outputs of a task. We create one
        artifact containing the agent's response.
        """
        artifacts: list[A2AArtifact] = []

        if result.response:
            artifacts.append(
                A2AArtifact(
                    name="response",
                    description="Agent response",
                    parts=[A2ATextPart(text=result.response)],
                    metadata={
                        "success": result.success,
                        "steps_taken": result.steps_taken,
                    },
                )
            )

        # If there were tool calls, add them as a data artifact
        if result.tool_calls:
            serialized_calls = [
                tc.model_dump(mode="json") if hasattr(tc, "model_dump") else tc
                for tc in result.tool_calls
            ]
            artifacts.append(
                A2AArtifact(
                    name="tool_calls",
                    description="Tools used during execution",
                    parts=[A2ADataPart(data={"tool_calls": serialized_calls})],
                )
            )

        return artifacts

    def _extract_message_text(self, history: list[A2AMessage]) -> str:
        """Extract text from the most recent user message."""
        for msg in reversed(history):
            if msg.role == "user":
                for part in msg.parts:
                    if isinstance(part, A2ATextPart):
                        return part.text

        # Fallback: just get any text from any message
        for msg in reversed(history):
            for part in msg.parts:
                if isinstance(part, A2ATextPart):
                    return part.text

        return ""

    def _create_error_message(self, error: str) -> A2AMessage:
        """Create an A2A error message."""
        return A2AMessage(
            role="agent",
            parts=[A2ATextPart(text=f"Error: {error}")],
        )


# =============================================================================
# A2A APP FACTORY
# =============================================================================


def get_albus_skills(runtime: Runtime) -> list[A2ASkill]:
    """Get Albus skills formatted for A2A AgentCard."""
    skills: list[A2ASkill] = []

    # Add skills from runtime
    for skill_id, skill in runtime.skills.items():
        skills.append(
            A2ASkill(
                id=skill_id,
                name=getattr(skill, "name", skill_id),
                description=getattr(skill, "description", f"Albus skill: {skill_id}"),
                tags=list(getattr(skill, "tags", [])),
                input_modes=["text/plain", "application/json"],
                output_modes=["text/plain", "application/json"],
            )
        )

    # Add a default "general" skill
    if not skills:
        skills.append(
            A2ASkill(
                id="general",
                name="General Assistant",
                description="General-purpose AI assistant with tool use capabilities",
                tags=["general", "assistant", "tools"],
                input_modes=["text/plain"],
                output_modes=["text/plain", "application/json"],
            )
        )

    return skills


def create_a2a_app(
    runtime: Runtime,
    *,
    name: str = "Albus",
    description: str = "Albus AgentOS - Event-driven AI agent runtime",
    version: str = "0.7.0",
    storage: Any | None = None,
    broker: Any | None = None,
) -> Any:
    """Create a FastA2A ASGI application for the Albus runtime.

    This creates an A2A-compliant server that exposes Albus agents
    via the A2A protocol, enabling interoperability with other
    agent frameworks.

    Args:
        runtime: Runtime instance
        name: Agent name for A2A card
        description: Agent description for A2A card
        version: Agent version for A2A card
        storage: Optional custom storage (defaults to InMemoryStorage)
        broker: Optional custom broker (defaults to InMemoryBroker)

    Returns
    -------
        FastA2A ASGI application

    Example:
        from server.runtime import Runtime
        from server.transport.a2a import create_a2a_app

        runtime = Runtime.create(...)
        a2a_app = create_a2a_app(runtime)

        # Run with uvicorn
        import uvicorn
        uvicorn.run(a2a_app, host="0.0.0.0", port=8081)
    """
    if not _FASTA2A_AVAILABLE:
        raise ImportError(
            "FastA2A is not installed. Install with: uv pip install albusos[a2a]"
        )

    # Use provided or create default storage/broker
    if storage is None:
        storage = InMemoryStorage()
    if broker is None:
        broker = InMemoryBroker()

    # Create worker
    worker = AlbusWorker(
        runtime=runtime,
        broker=broker,
        storage=storage,
        default_agent_id="host",
    )

    # Build skills for agent card (convert to dicts for fasta2a)
    skills = [s.model_dump(mode="json") for s in get_albus_skills(runtime)]

    # Lifespan context manager
    @asynccontextmanager
    async def lifespan(app: FastA2A) -> AsyncIterator[None]:
        # Start the task manager and worker
        async with app.task_manager, worker.run():
            logger.info("[A2A] Albus A2A server started")
            yield
            logger.info("[A2A] Albus A2A server stopping")

    # Create FastA2A app
    app = FastA2A(
        storage=storage,
        broker=broker,
        lifespan=lifespan,
        # Agent card configuration
        name=name,
        description=description,
        version=version,
        skills=skills,
    )

    return app


# =============================================================================
# UVICORN RUNNER
# =============================================================================


async def run_a2a_server(
    runtime: Runtime,
    *,
    host: str = "127.0.0.1",
    port: int = 8081,
    **kwargs: Any,
) -> None:
    """Run the A2A server with uvicorn.

    This is a convenience function for running the A2A server
    as a standalone process.

    Args:
        runtime: Runtime instance
        host: Host to bind to
        port: Port to bind to (default: 8081, different from main server)
        **kwargs: Additional arguments for create_a2a_app
    """
    if not _FASTA2A_AVAILABLE:
        raise ImportError(
            "FastA2A is not installed. Install with: uv pip install albusos[a2a]"
        )

    try:
        import uvicorn  # noqa: PLC0415
    except ImportError:
        raise ImportError(
            "uvicorn is not installed. Install with: uv pip install albusos[a2a]"
        )

    app = create_a2a_app(runtime, **kwargs)

    config = uvicorn.Config(
        app,
        host=host,
        port=port,
        log_level="info",
    )
    server = uvicorn.Server(config)

    logger.info(f"[A2A] Starting A2A server at http://{host}:{port}")
    print(f"\n🔗 A2A Protocol server at http://{host}:{port}")
    print("   Agent card: GET /")
    print("   JSON-RPC:   POST /\n")

    await server.serve()


__all__ = [
    "A2AArtifact",
    "A2ADataPart",
    "A2AMessage",
    "A2ASkill",
    "A2ATextPart",
    "AlbusWorker",
    "create_a2a_app",
    "get_albus_skills",
    "is_a2a_available",
    "run_a2a_server",
]
