import asyncio
from typing import Union, Sequence, Optional, List, Dict, AsyncGenerator
import json
import httpx

from acp_sdk.client import Client as AcpSdkClient
from acp_sdk.models import Event, Message
from acp_sdk.models.models import MessagePart

MessageInput = Union[str, dict, tuple]


class AcpAgentResolutionError(Exception):
    pass


def _run_sync(coro):
    """
    Safely run an async coroutine from sync code.

    asyncio.run() raises RuntimeError if called while an event loop is
    already running (e.g. inside FastAPI/Jupyter/another async context).
    This helper detects that case and gives a clear error instead of a
    confusing crash, while still working normally in plain scripts.
    """
    try:
        asyncio.get_running_loop()
    except RuntimeError:
        # No running loop -> safe to use asyncio.run()
        return asyncio.run(coro)
    else:
        raise RuntimeError(
            "This sync method was called from inside a running event loop "
            "(e.g. FastAPI, Jupyter, or another async function). "
            "Use the async version instead (e.g. 'await client.chat_async(...)' "
            "or 'await client.list_agents_async(...)')."
        )


class AcpClient:
    def __init__(self, base_url: str, agent: Optional[str] = None, timeout: float = 60.0):
        self.base_url = base_url.rstrip("/")
        self.agent = agent
        self.timeout = timeout
        self.headers = {"Content-Type": "application/json"}

    async def list_agents_async(self) -> List[Dict[str, str]]:
        agents = []
        async with AcpSdkClient(base_url=self.base_url, headers=self.headers) as client:
            async for agent in client.agents():
                agents.append({
                    "name": agent.name,
                    "description": agent.description or "No description provided.",
                    "input_content_types": str(agent.input_content_types),
                    "output_content_types": str(agent.output_content_types),
                })
        return agents

    def list_agents(self) -> List[Dict[str, str]]:
        return _run_sync(self.list_agents_async())

    async def _resolve_agent(self, override: Optional[str]) -> str:
        if override:
            return override
        if self.agent:
            return self.agent

        available = await self.list_agents_async()

        if not available:
            raise AcpAgentResolutionError(
                "No agent found on this server. Check that the server is running "
                "and reachable at " + self.base_url
            )
        if len(available) == 1:
            return available[0]["name"]

        names = [a["name"] for a in available]
        details = "\n".join(f"  - {a['name']}: {a['description']}" for a in available)
        raise AcpAgentResolutionError(
            f"Multiple agents are available on this server and none was specified:\n"
            f"{details}\n\n"
            f"Fix this by passing agent=... to chat(), e.g. agent='{names[0]}', "
            f"or set a default with AcpClient(base_url=..., agent='{names[0]}')."
        )

    @staticmethod
    def _normalize(messages: List[Dict[str, str]]) -> List[Dict[str, str]]:
        if not isinstance(messages, list):
            raise AcpAgentResolutionError("Message must be List of Dict.")

        normalized: List[Dict[str, str]] = []

        for i, message in enumerate(messages):
            if not isinstance(message, dict):
                raise AcpAgentResolutionError(
                    f"Message at index {i} must be dictionary, got {type(message).__name__}."
                )

            if "role" not in message:
                raise AcpAgentResolutionError(f"Message at index {i} is missing the 'role' field.")
            if "content" not in message:
                raise AcpAgentResolutionError(f"Message at index {i} is missing the 'content' field.")
            if not isinstance(message["role"], str):
                raise AcpAgentResolutionError(f"'role' at index {i} must be a string.")
            if not isinstance(message["content"], str):
                raise AcpAgentResolutionError(f"'content' at index {i} must be a string.")

            normalized.append({"role": message["role"], "content": message["content"]})

        return normalized

    @staticmethod
    def _to_payload(agent: str, canonical: List[Dict[str, str]], model: str) -> dict:
        input_messages = [
            {
                "role": m["role"],
                "parts": [{"content": m["content"], "content_type": "text/plain"}],
            }
            for m in canonical
        ]

        # Append model metadata as an extra part on the LAST message only
        if input_messages and model:
            input_messages[-1]["parts"].append({
                "content": json.dumps({"model": model}),
                "content_type": "application/json",
            })
        return {"agent_name": agent, "input": input_messages}

    @staticmethod
    def _to_acp_messages(canonical: List[Dict[str, str]], model: str) -> List[Message]:
        """Build acp_sdk Message objects (used by the SDK's native streaming
        client), embedding the model as a JSON part on the last message —
        same convention as _to_payload uses for the raw HTTP call."""
        messages = [
            Message(role=m["role"], parts=[MessagePart(content=m["content"], content_type="text/plain")])
            for m in canonical
        ]

        if messages and model:
            messages[-1].parts.append(
                MessagePart(content=json.dumps({"model": model}), content_type="application/json")
            )
        return messages

    @staticmethod
    def _extract_text(data: dict) -> str:
        """
        Concatenate text/plain content across EVERY message in `output`,
        not just output[0].

        Why: your server yields a brand-new Message per streamed token,
        so a multi-token response arrives as many separate entries in
        `output` rather than one message with many parts. Reading only
        output[0] used to mean chat() showed just the first token.
        """
        output = data.get("output", [])
        if not output:
            return ""

        text = []
        for message in output:
            for part in message.get("parts", []):
                if part.get("content_type") == "text/plain":
                    text.append(part.get("content", ""))
        return "".join(text)

    async def _run_async(self, messages: List[Dict[str, str]], model: str, agent: Optional[str] = None) -> str:
        resolved_agent = await self._resolve_agent(agent)
        canonical = self._normalize(messages)
        payload = self._to_payload(resolved_agent, canonical, model)

        async with httpx.AsyncClient(base_url=self.base_url, timeout=self.timeout) as client:
            resp = await client.post("/runs", json=payload, headers=self.headers)
            if resp.status_code == 422:
                raise AcpAgentResolutionError(f"Validation error from server: {resp.json()}")
            resp.raise_for_status()
            return self._extract_text(resp.json())

    def chat(self, messages: Sequence[MessageInput], model: str, agent: Optional[str] = None) -> str:
        """Multi-turn call (sync). Pass a list of {'role': ..., 'content': ...} dicts.
        Do not call this from inside a running event loop — use chat_async instead."""
        return _run_sync(self._run_async(messages, model, agent))

    async def chat_async(self, messages: List[Dict[str, str]], model: str, agent: Optional[str] = None) -> str:
        """Multi-turn call (async). Use this inside FastAPI, Jupyter, or any
        other code that already has an event loop running."""
        return await self._run_async(messages, model, agent)

    async def chat_stream(
        self, messages: List[Dict[str, str]], model: str, agent: Optional[str] = None
    ) -> AsyncGenerator[Event, None]:
        """
        Multi-turn streaming call. Yields the raw ACP `Event` objects exactly
        as they come off the server's SSE stream (RunCreatedEvent,
        MessagePartEvent, RunCompletedEvent, ErrorEvent, etc.) — nothing is
        parsed, flattened, or extracted, so the caller gets the whole object
        and can pull whatever fields it needs.

        Usage:
            async for event in client.chat_stream(messages, model="gpt-4o"):
                print(event)
        """
        resolved_agent = await self._resolve_agent(agent)
        canonical = self._normalize(messages)
        acp_messages = self._to_acp_messages(canonical, model)

        async with AcpSdkClient(base_url=self.base_url, headers=self.headers) as client:
            async for event in client.run_stream(acp_messages, agent=resolved_agent):
                yield event