from __future__ import annotations

import json
import time
import logging
from dataclasses import dataclass, field
from typing import Dict, Optional, Any, List, Iterator, AsyncIterator

import httpx

logger = logging.getLogger("promptev")

# ── Data classes ─────────────────────────────────────────


@dataclass
class AgentSession:
    """Returned by start_agent / astart_agent."""

    session_token: str
    chatbot_id: str
    name: str
    memory_enabled: bool
    messages: List[Dict[str, Any]] = field(default_factory=list)


@dataclass
class AgentEvent:
    """A single SSE event from the agent stream."""

    type: str  # "thoughts", "processing", "done", "error", "approval_required"
    output: str
    raw: Dict[str, Any] = field(default_factory=dict)


# ── Exceptions ───────────────────────────────────────────


class PromptevError(Exception):
    """Base exception for all SDK errors."""

    def __init__(
        self,
        message: str,
        status_code: Optional[int] = None,
        response_text: Optional[str] = None,
    ):
        super().__init__(message)
        self.status_code = status_code
        self.response_text = response_text


class ValidationError(PromptevError):
    """400 — missing variables, bad input."""


class AuthenticationError(PromptevError):
    """401/403 — invalid API key, inactive agent."""


class NotFoundError(PromptevError):
    """404 — project, prompt, or agent not found."""


class RateLimitError(PromptevError):
    """429 — API quota exceeded."""


class ServerError(PromptevError):
    """5xx — server error."""


class NetworkError(PromptevError):
    """Connection failed, timeout, DNS error."""


# ── Client ───────────────────────────────────────────────

_RETRYABLE_STATUS = {502, 503, 504}
_DEFAULT_MAX_RETRIES = 2
_DEFAULT_BACKOFF = 0.5  # seconds, doubles each retry


class PromptevClient:
    """
    Python SDK for Promptev.

    Supports:
      - Prompt compilation  (get_prompt / aget_prompt)
      - Agent sessions      (start_agent / astart_agent)
      - Agent streaming     (stream_agent / astream_agent)
    """

    def __init__(
        self,
        project_key: str,
        base_url: str = "https://api.promptev.ai",
        headers: Optional[Dict[str, str]] = None,
        timeout: float = 30.0,
        max_retries: int = _DEFAULT_MAX_RETRIES,
    ) -> None:
        self.project_key = project_key
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.max_retries = max_retries

        merged_headers = {"Content-Type": "application/json", **(headers or {})}

        self._client = httpx.Client(
            base_url=self.base_url,
            headers=merged_headers,
            timeout=self.timeout,
        )
        self._aclient = httpx.AsyncClient(
            base_url=self.base_url,
            headers=merged_headers,
            timeout=self.timeout,
        )

    # ── Resource management ──────────────────────────────

    def close(self) -> None:
        """Close the underlying HTTP clients (sync)."""
        self._client.close()

    async def aclose(self) -> None:
        """Close both HTTP clients (async)."""
        await self._aclient.aclose()
        self._client.close()

    def __enter__(self) -> PromptevClient:
        return self

    def __exit__(self, *exc: Any) -> None:
        self.close()

    async def __aenter__(self) -> PromptevClient:
        return self

    async def __aexit__(self, *exc: Any) -> None:
        await self.aclose()

    # ── Error handling ───────────────────────────────────

    @staticmethod
    def _raise_for_status(resp: httpx.Response) -> None:
        """Map HTTP status codes to typed SDK exceptions."""
        if resp.status_code < 400:
            return

        try:
            detail = resp.json().get("detail", resp.text)
        except Exception:
            detail = resp.text or "Unknown error"

        if not isinstance(detail, str):
            detail = str(detail)

        code = resp.status_code
        text = resp.text

        if code == 400:
            raise ValidationError(detail, code, text)
        if code == 401:
            raise AuthenticationError(detail, code, text)
        if code == 403:
            raise AuthenticationError(detail, code, text)
        if code == 404:
            raise NotFoundError(detail, code, text)
        if code == 429:
            raise RateLimitError(detail, code, text)
        if code >= 500:
            raise ServerError(detail, code, text)
        raise PromptevError(detail, code, text)

    # ── Retry helpers ────────────────────────────────────

    def _request_with_retry(
        self, method: str, url: str, **kwargs: Any
    ) -> httpx.Response:
        """Sync request with automatic retry on transient failures."""
        last_exc: Optional[Exception] = None

        for attempt in range(self.max_retries + 1):
            try:
                resp = self._client.request(method, url, **kwargs)
                if resp.status_code not in _RETRYABLE_STATUS or attempt == self.max_retries:
                    return resp
                last_exc = ServerError(
                    f"Server error ({resp.status_code})", resp.status_code
                )
            except httpx.HTTPError as e:
                last_exc = NetworkError(f"Network error: {e}")
                if attempt == self.max_retries:
                    raise last_exc from e

            delay = _DEFAULT_BACKOFF * (2 ** attempt)
            logger.debug("Retry %d/%d after %.1fs", attempt + 1, self.max_retries, delay)
            time.sleep(delay)

        raise last_exc  # type: ignore[misc]

    async def _arequest_with_retry(
        self, method: str, url: str, **kwargs: Any
    ) -> httpx.Response:
        """Async request with automatic retry on transient failures."""
        import asyncio

        last_exc: Optional[Exception] = None

        for attempt in range(self.max_retries + 1):
            try:
                resp = await self._aclient.request(method, url, **kwargs)
                if resp.status_code not in _RETRYABLE_STATUS or attempt == self.max_retries:
                    return resp
                last_exc = ServerError(
                    f"Server error ({resp.status_code})", resp.status_code
                )
            except httpx.HTTPError as e:
                last_exc = NetworkError(f"Network error: {e}")
                if attempt == self.max_retries:
                    raise last_exc from e

            delay = _DEFAULT_BACKOFF * (2 ** attempt)
            logger.debug("Retry %d/%d after %.1fs", attempt + 1, self.max_retries, delay)
            await asyncio.sleep(delay)

        raise last_exc  # type: ignore[misc]

    # ── Prompt API ───────────────────────────────────────

    def run_prompt(
        self,
        prompt_key: str,
        query: str,
        variables: Optional[Dict[str, str]] = None,
    ) -> str:
        """
        Compile and execute a prompt (sync).

        If the prompt has a model configured, it compiles the template,
        sends it to the LLM, and returns the AI response. If no model
        is configured, it returns the compiled template string.

        Args:
            prompt_key: The prompt name as configured in Promptev.
            query: The user's question or input for the LLM.
            variables: Optional dict of template variables.

        Returns:
            The AI response (if model configured) or compiled prompt string.

        Raises:
            ValidationError: Missing required variables or bad input (400).
            NotFoundError: Prompt or project not found (404).
        """
        url = f"/api/sdk/v1/prompt/client/{self.project_key}/{prompt_key}"
        merged = {"query": query, **(variables or {})}
        resp = self._request_with_retry("POST", url, json={"variables": merged})
        self._raise_for_status(resp)

        data = resp.json()
        compiled = data.get("prompt")
        if not isinstance(compiled, str):
            raise PromptevError("Unexpected response: missing 'prompt' field")
        return compiled

    async def arun_prompt(
        self,
        prompt_key: str,
        query: str,
        variables: Optional[Dict[str, str]] = None,
    ) -> str:
        """
        Compile and execute a prompt (async).

        If the prompt has a model configured, it compiles the template,
        sends it to the LLM, and returns the AI response. If no model
        is configured, it returns the compiled template string.

        Args:
            prompt_key: The prompt name as configured in Promptev.
            query: The user's question or input for the LLM.
            variables: Optional dict of template variables.

        Returns:
            The AI response (if model configured) or compiled prompt string.

        Raises:
            ValidationError: Missing required variables or bad input (400).
            NotFoundError: Prompt or project not found (404).
        """
        url = f"/api/sdk/v1/prompt/client/{self.project_key}/{prompt_key}"
        merged = {"query": query, **(variables or {})}
        resp = await self._arequest_with_retry(
            "POST", url, json={"variables": merged}
        )
        self._raise_for_status(resp)

        data = resp.json()
        compiled = data.get("prompt")
        if not isinstance(compiled, str):
            raise PromptevError("Unexpected response: missing 'prompt' field")
        return compiled

    def stream_prompt(
        self,
        prompt_key: str,
        query: str,
        variables: Optional[Dict[str, str]] = None,
    ) -> Iterator[AgentEvent]:
        """
        Compile and execute a prompt with streaming (sync generator).

        Use this when the prompt has tools attached or you want real-time
        output. Returns the same SSE event types as agent streaming.

        Args:
            prompt_key: The prompt name as configured in Promptev.
            query: The user's question or input.
            variables: Optional dict of template variables.

        Yields:
            AgentEvent for each SSE event (thoughts, processing, done, error).
        """
        url = f"/api/sdk/v1/prompt/client/{self.project_key}/{prompt_key}"
        merged = {"query": query, **(variables or {})}
        req = self._client.build_request(
            "POST",
            url,
            json={"variables": merged},
            params={"stream": "true"},
        )
        resp = self._client.send(req, stream=True)
        try:
            if resp.status_code >= 400:
                resp.read()
                self._raise_for_status(resp)

            content_type = resp.headers.get("content-type", "")

            # Non-streaming fallback: prompt has no model, backend returns JSON
            if "text/event-stream" not in content_type:
                resp.read()
                data = resp.json()
                yield AgentEvent(
                    type="done",
                    output=data.get("prompt", ""),
                    raw=data,
                )
                return

            for line in resp.iter_lines():
                if not line.startswith("data: "):
                    continue
                try:
                    event_data = json.loads(line[6:])
                except json.JSONDecodeError:
                    continue

                event = AgentEvent(
                    type=event_data.get("type", "unknown"),
                    output=event_data.get("output", ""),
                    raw=event_data,
                )
                yield event

                if event.type in ("done", "error"):
                    break
        finally:
            resp.close()

    async def astream_prompt(
        self,
        prompt_key: str,
        query: str,
        variables: Optional[Dict[str, str]] = None,
    ) -> AsyncIterator[AgentEvent]:
        """
        Compile and execute a prompt with streaming (async generator).

        Use this when the prompt has tools attached or you want real-time
        output. Returns the same SSE event types as agent streaming.

        Args:
            prompt_key: The prompt name as configured in Promptev.
            query: The user's question or input.
            variables: Optional dict of template variables.

        Yields:
            AgentEvent for each SSE event (thoughts, processing, done, error).
        """
        url = f"/api/sdk/v1/prompt/client/{self.project_key}/{prompt_key}"
        merged = {"query": query, **(variables or {})}
        req = self._aclient.build_request(
            "POST",
            url,
            json={"variables": merged},
            params={"stream": "true"},
        )
        resp = await self._aclient.send(req, stream=True)
        try:
            if resp.status_code >= 400:
                await resp.aread()
                self._raise_for_status(resp)

            content_type = resp.headers.get("content-type", "")

            # Non-streaming fallback: prompt has no model, backend returns JSON
            if "text/event-stream" not in content_type:
                await resp.aread()
                data = resp.json()
                yield AgentEvent(
                    type="done",
                    output=data.get("prompt", ""),
                    raw=data,
                )
                return

            async for line in resp.aiter_lines():
                if not line.startswith("data: "):
                    continue
                try:
                    event_data = json.loads(line[6:])
                except json.JSONDecodeError:
                    continue

                event = AgentEvent(
                    type=event_data.get("type", "unknown"),
                    output=event_data.get("output", ""),
                    raw=event_data,
                )
                yield event

                if event.type in ("done", "error"):
                    break
        finally:
            await resp.aclose()

    # ── Agent API ────────────────────────────────────────

    def start_agent(
        self,
        chatbot_id: str,
        *,
        visitor: Optional[str] = None,
        platform: str = "sdk",
    ) -> AgentSession:
        """
        Start a new agent chat session (sync).

        Args:
            chatbot_id: The agent/chatbot UUID.
            visitor: Display name for the user (defaults to "anonymous").
            platform: Platform identifier (defaults to "sdk").

        Returns:
            AgentSession with session_token for streaming.

        Raises:
            NotFoundError: Agent not found or not in this project (404).
            AuthenticationError: Agent is not active (403).
        """
        url = f"/api/sdk/v1/agent/{self.project_key}/{chatbot_id}/start"
        payload: Dict[str, Any] = {}
        if visitor is not None:
            payload["visitor"] = visitor
        if platform != "sdk":
            payload["platform"] = platform

        resp = self._request_with_retry("POST", url, json=payload)
        self._raise_for_status(resp)

        data = resp.json()
        return AgentSession(
            session_token=data["session_token"],
            chatbot_id=data["chatbot_id"],
            name=data["name"],
            memory_enabled=data.get("memory_enabled", False),
            messages=data.get("messages", []),
        )

    async def astart_agent(
        self,
        chatbot_id: str,
        *,
        visitor: Optional[str] = None,
        platform: str = "sdk",
    ) -> AgentSession:
        """
        Start a new agent chat session (async).

        Args:
            chatbot_id: The agent/chatbot UUID.
            visitor: Display name for the user (defaults to "anonymous").
            platform: Platform identifier (defaults to "sdk").

        Returns:
            AgentSession with session_token for streaming.

        Raises:
            NotFoundError: Agent not found or not in this project (404).
            AuthenticationError: Agent is not active (403).
        """
        url = f"/api/sdk/v1/agent/{self.project_key}/{chatbot_id}/start"
        payload: Dict[str, Any] = {}
        if visitor is not None:
            payload["visitor"] = visitor
        if platform != "sdk":
            payload["platform"] = platform

        resp = await self._arequest_with_retry("POST", url, json=payload)
        self._raise_for_status(resp)

        data = resp.json()
        return AgentSession(
            session_token=data["session_token"],
            chatbot_id=data["chatbot_id"],
            name=data["name"],
            memory_enabled=data.get("memory_enabled", False),
            messages=data.get("messages", []),
        )

    def stream_agent(
        self,
        chatbot_id: str,
        *,
        session_token: str,
        query: str,
    ) -> Iterator[AgentEvent]:
        """
        Stream an agent response as SSE events (sync generator).

        Args:
            chatbot_id: The agent/chatbot UUID.
            session_token: Token from start_agent().
            query: The user's message.

        Yields:
            AgentEvent for each SSE event (thoughts, processing, done, error).
        """
        url = f"/api/sdk/v1/agent/{self.project_key}/{chatbot_id}/stream"
        req = self._client.build_request(
            "POST",
            url,
            json={"query": query},
            params={"session_token": session_token},
        )
        resp = self._client.send(req, stream=True)
        try:
            if resp.status_code >= 400:
                resp.read()
                self._raise_for_status(resp)

            for line in resp.iter_lines():
                if not line.startswith("data: "):
                    continue
                try:
                    event_data = json.loads(line[6:])
                except json.JSONDecodeError:
                    continue

                event = AgentEvent(
                    type=event_data.get("type", "unknown"),
                    output=event_data.get("output", ""),
                    raw=event_data,
                )
                yield event

                if event.type in ("done", "error"):
                    break
        finally:
            resp.close()

    async def astream_agent(
        self,
        chatbot_id: str,
        *,
        session_token: str,
        query: str,
    ) -> AsyncIterator[AgentEvent]:
        """
        Stream an agent response as SSE events (async generator).

        Args:
            chatbot_id: The agent/chatbot UUID.
            session_token: Token from start_agent().
            query: The user's message.

        Yields:
            AgentEvent for each SSE event (thoughts, processing, done, error).
        """
        url = f"/api/sdk/v1/agent/{self.project_key}/{chatbot_id}/stream"
        req = self._aclient.build_request(
            "POST",
            url,
            json={"query": query},
            params={"session_token": session_token},
        )
        resp = await self._aclient.send(req, stream=True)
        try:
            if resp.status_code >= 400:
                await resp.aread()
                self._raise_for_status(resp)

            async for line in resp.aiter_lines():
                if not line.startswith("data: "):
                    continue
                try:
                    event_data = json.loads(line[6:])
                except json.JSONDecodeError:
                    continue

                event = AgentEvent(
                    type=event_data.get("type", "unknown"),
                    output=event_data.get("output", ""),
                    raw=event_data,
                )
                yield event

                if event.type in ("done", "error"):
                    break
        finally:
            await resp.aclose()
