"""a21e API client — sync and async."""

from __future__ import annotations

import json
from typing import Any, Generator, AsyncGenerator, Optional

import httpx


class A21eError(Exception):
    """Error from the a21e API."""

    def __init__(self, message: str, status: int = 0, code: str | None = None, data: dict | None = None):
        super().__init__(message)
        self.status = status
        self.code = code
        self.data = data or {}


class A21eClient:
    """Synchronous client for the a21e API.

    Usage::

        client = A21eClient(api_key="a21e_...")
        result = client.run("prompt-id", "Hello world")
        print(result["output"])
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = "https://api.a21e.com",
        timeout: float = 60.0,
    ):
        if not api_key:
            raise ValueError("api_key is required")
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self._http = httpx.Client(
            base_url=self.base_url,
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            timeout=timeout,
        )

    def close(self) -> None:
        self._http.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()

    def _request(self, method: str, path: str, **kwargs) -> Any:
        res = self._http.request(method, path, **kwargs)
        if res.status_code >= 400:
            body = res.json() if res.headers.get("content-type", "").startswith("application/json") else {}
            extra = {k: v for k, v in body.items() if k not in ("error", "message", "code")}
            raise A21eError(
                body.get("error", body.get("message", f"API error: {res.status_code}")),
                status=res.status_code,
                code=body.get("code"),
                data=extra,
            )
        return res.json()

    # ── Prompts ──────────────────────────────────────────────────────────

    def list_prompts(
        self, *, topic: str | None = None, q: str | None = None
    ) -> list[dict]:
        params = {}
        if topic:
            params["topic"] = topic
        if q:
            params["q"] = q
        return self._request("GET", "/v1/prompts", params=params)

    def get_prompt(self, prompt_id: str) -> dict:
        return self._request("GET", f"/v1/prompts/{prompt_id}")

    def list_versions(self, prompt_id: str) -> list[dict]:
        return self._request("GET", f"/v1/prompts/{prompt_id}/versions")

    def run(
        self,
        prompt_id: str,
        input: str,
        *,
        model: str | None = None,
        version: int | None = None,
        provider_api_key: str | None = None,
    ) -> dict:
        """Run a prompt and return the result."""
        body: dict[str, Any] = {"input": input}
        if model:
            body["model"] = model
        if version is not None:
            body["version"] = version
        if provider_api_key:
            body["provider_api_key"] = provider_api_key
        return self._request("POST", f"/v1/prompts/{prompt_id}/run", json=body)

    def run_stream(
        self,
        prompt_id: str,
        input: str,
        *,
        model: str | None = None,
        version: int | None = None,
        provider_api_key: str | None = None,
    ) -> Generator[str, None, None]:
        """Run a prompt with SSE streaming. Yields text chunks."""
        body: dict[str, Any] = {"input": input, "stream": True}
        if model:
            body["model"] = model
        if version is not None:
            body["version"] = version
        if provider_api_key:
            body["provider_api_key"] = provider_api_key

        with self._http.stream(
            "POST", f"/v1/prompts/{prompt_id}/run", json=body
        ) as res:
            if res.status_code >= 400:
                res.read()
                data = res.json() if res.headers.get("content-type", "").startswith("application/json") else {}
                raise A21eError(
                    data.get("error", f"API error: {res.status_code}"),
                    status=res.status_code,
                    code=data.get("code"),
                )
            for line in res.iter_lines():
                if line.startswith("data: "):
                    data = line[6:]
                    if data == "[DONE]":
                        return
                    try:
                        parsed = json.loads(data)
                        content = parsed.get("choices", [{}])[0].get("delta", {}).get("content")
                        if content:
                            yield content
                    except json.JSONDecodeError:
                        continue

    # ── Topics ───────────────────────────────────────────────────────────

    def list_topics(self) -> list[dict]:
        return self._request("GET", "/v1/topics")

    # ── Library ──────────────────────────────────────────────────────────

    def list_library(self) -> list[dict]:
        return self._request("GET", "/v1/library")

    def add_to_library(
        self, prompt_id: str, prompt_version_id: str | None = None
    ) -> dict:
        body: dict[str, Any] = {"prompt_id": prompt_id}
        if prompt_version_id:
            body["prompt_version_id"] = prompt_version_id
        return self._request("POST", "/v1/library", json=body)

    def update_library_entry(
        self, prompt_id: str, prompt_version_id: str
    ) -> dict:
        return self._request(
            "PATCH",
            f"/v1/library/{prompt_id}",
            json={"prompt_version_id": prompt_version_id},
        )

    # ── Workspace / Favorites ───────────────────────────────────────────

    def get_workspace(self) -> dict:
        return self._request("GET", "/v1/workspace")

    def list_favorites(self) -> list[dict]:
        return self._request("GET", "/v1/workspace/favorites")

    def add_favorite(
        self, prompt_id: str, prompt_version_id: str | None = None
    ) -> dict:
        body: dict[str, Any] = {"prompt_id": prompt_id}
        if prompt_version_id:
            body["prompt_version_id"] = prompt_version_id
        return self._request("POST", "/v1/workspace/favorites", json=body)

    def update_favorite_entry(
        self, prompt_id: str, prompt_version_id: str
    ) -> dict:
        return self._request(
            "PATCH",
            f"/v1/workspace/favorites/{prompt_id}",
            json={"prompt_version_id": prompt_version_id},
        )

    # ── Intents ──────────────────────────────────────────────────────────

    def submit_intent(self, intent: dict) -> dict:
        return self._request("POST", "/v1/intents", json=intent)

    def list_intents(
        self,
        *,
        limit: int | None = None,
        offset: int | None = None,
        status: str | None = None,
    ) -> dict:
        params = {}
        if limit is not None:
            params["limit"] = str(limit)
        if offset is not None:
            params["offset"] = str(offset)
        if status:
            params["status"] = status
        return self._request("GET", "/v1/intents", params=params)

    def get_intent(self, intent_id: str) -> dict:
        return self._request("GET", f"/v1/intents/{intent_id}")

    def clarify_intent(self, intent_id: str, clarification_input: str) -> dict:
        return self._request(
            "POST",
            f"/v1/intents/{intent_id}/clarify",
            json={"clarification_input": clarification_input},
        )

    # ── Packages ─────────────────────────────────────────────────────────

    def execute_package(
        self,
        package_id: str,
        *,
        input: str | None = None,
        model: str | None = None,
        provider_api_key: str | None = None,
    ) -> dict:
        body: dict[str, Any] = {}
        if input:
            body["input"] = input
        if model:
            body["model"] = model
        if provider_api_key:
            body["provider_api_key"] = provider_api_key
        return self._request(
            "POST", f"/v1/packages/{package_id}/execute", json=body
        )

    # ── Notifications ────────────────────────────────────────────────────

    def get_notifications(self, *, limit: int | None = None) -> dict:
        params = {}
        if limit is not None:
            params["limit"] = str(limit)
        return self._request("GET", "/v1/notifications", params=params)

    def mark_notifications_read(self, ids: list[str] | str) -> dict:
        return self._request("POST", "/v1/notifications/read", json={"ids": ids})

    # ── Analytics ────────────────────────────────────────────────────────

    def get_creator_analytics(
        self,
        *,
        period_start: str | None = None,
        period_end: str | None = None,
    ) -> dict:
        params = {}
        if period_start:
            params["period_start"] = period_start
        if period_end:
            params["period_end"] = period_end
        return self._request("GET", "/v1/analytics/creator", params=params)

    # ── Usage ────────────────────────────────────────────────────────────

    def get_usage_summary(self, *, days: int | None = None) -> dict:
        params = {}
        if days is not None:
            params["days"] = str(days)
        return self._request("GET", "/v1/usage/summary", params=params)

    def get_usage_by_day(self, *, days: int | None = None) -> dict:
        params = {}
        if days is not None:
            params["days"] = str(days)
        return self._request("GET", "/v1/usage/by-day", params=params)

    # ── Feedback ─────────────────────────────────────────────────────────

    def submit_feedback(self, feedback: dict) -> dict:
        return self._request("POST", "/v1/feedback", json=feedback)

    # ── RPC ──────────────────────────────────────────────────────────────

    def rpc(self, method: str, params: dict) -> Any:
        """Send a raw JSON-RPC 2.0 request."""
        import time

        response = self._request(
            "POST",
            "/v1/rpc",
            json={
                "jsonrpc": "2.0",
                "method": method,
                "params": params,
                "id": int(time.time() * 1000),
            },
        )
        if "error" in response:
            raise A21eError(
                response["error"].get("message", "RPC error"),
                code=response["error"].get("code"),
            )
        return response.get("result")


class AsyncA21eClient:
    """Async client for the a21e API.

    Usage::

        async with AsyncA21eClient(api_key="a21e_...") as client:
            result = await client.run("prompt-id", "Hello world")
            print(result["output"])
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = "https://api.a21e.com",
        timeout: float = 60.0,
    ):
        if not api_key:
            raise ValueError("api_key is required")
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self._http = httpx.AsyncClient(
            base_url=self.base_url,
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            timeout=timeout,
        )

    async def close(self) -> None:
        await self._http.aclose()

    async def __aenter__(self):
        return self

    async def __aexit__(self, *args):
        await self.close()

    async def _request(self, method: str, path: str, **kwargs) -> Any:
        res = await self._http.request(method, path, **kwargs)
        if res.status_code >= 400:
            body = res.json() if res.headers.get("content-type", "").startswith("application/json") else {}
            extra = {k: v for k, v in body.items() if k not in ("error", "message", "code")}
            raise A21eError(
                body.get("error", body.get("message", f"API error: {res.status_code}")),
                status=res.status_code,
                code=body.get("code"),
                data=extra,
            )
        return res.json()

    # ── Prompts ──────────────────────────────────────────────────────────

    async def list_prompts(
        self, *, topic: str | None = None, q: str | None = None
    ) -> list[dict]:
        params = {}
        if topic:
            params["topic"] = topic
        if q:
            params["q"] = q
        return await self._request("GET", "/v1/prompts", params=params)

    async def get_prompt(self, prompt_id: str) -> dict:
        return await self._request("GET", f"/v1/prompts/{prompt_id}")

    async def list_versions(self, prompt_id: str) -> list[dict]:
        return await self._request("GET", f"/v1/prompts/{prompt_id}/versions")

    async def run(
        self,
        prompt_id: str,
        input: str,
        *,
        model: str | None = None,
        version: int | None = None,
        provider_api_key: str | None = None,
    ) -> dict:
        body: dict[str, Any] = {"input": input}
        if model:
            body["model"] = model
        if version is not None:
            body["version"] = version
        if provider_api_key:
            body["provider_api_key"] = provider_api_key
        return await self._request(
            "POST", f"/v1/prompts/{prompt_id}/run", json=body
        )

    async def run_stream(
        self,
        prompt_id: str,
        input: str,
        *,
        model: str | None = None,
        version: int | None = None,
        provider_api_key: str | None = None,
    ) -> AsyncGenerator[str, None]:
        body: dict[str, Any] = {"input": input, "stream": True}
        if model:
            body["model"] = model
        if version is not None:
            body["version"] = version
        if provider_api_key:
            body["provider_api_key"] = provider_api_key

        async with self._http.stream(
            "POST", f"/v1/prompts/{prompt_id}/run", json=body
        ) as res:
            if res.status_code >= 400:
                await res.aread()
                data = res.json() if res.headers.get("content-type", "").startswith("application/json") else {}
                raise A21eError(
                    data.get("error", f"API error: {res.status_code}"),
                    status=res.status_code,
                    code=data.get("code"),
                )
            async for line in res.aiter_lines():
                if line.startswith("data: "):
                    data = line[6:]
                    if data == "[DONE]":
                        return
                    try:
                        parsed = json.loads(data)
                        content = (
                            parsed.get("choices", [{}])[0]
                            .get("delta", {})
                            .get("content")
                        )
                        if content:
                            yield content
                    except json.JSONDecodeError:
                        continue

    # ── Topics ───────────────────────────────────────────────────────────

    async def list_topics(self) -> list[dict]:
        return await self._request("GET", "/v1/topics")

    # ── Library ──────────────────────────────────────────────────────────

    async def list_library(self) -> list[dict]:
        return await self._request("GET", "/v1/library")

    async def add_to_library(
        self, prompt_id: str, prompt_version_id: str | None = None
    ) -> dict:
        body: dict[str, Any] = {"prompt_id": prompt_id}
        if prompt_version_id:
            body["prompt_version_id"] = prompt_version_id
        return await self._request("POST", "/v1/library", json=body)

    async def update_library_entry(
        self, prompt_id: str, prompt_version_id: str
    ) -> dict:
        return await self._request(
            "PATCH",
            f"/v1/library/{prompt_id}",
            json={"prompt_version_id": prompt_version_id},
        )

    # ── Workspace / Favorites ───────────────────────────────────────────

    async def get_workspace(self) -> dict:
        return await self._request("GET", "/v1/workspace")

    async def list_favorites(self) -> list[dict]:
        return await self._request("GET", "/v1/workspace/favorites")

    async def add_favorite(
        self, prompt_id: str, prompt_version_id: str | None = None
    ) -> dict:
        body: dict[str, Any] = {"prompt_id": prompt_id}
        if prompt_version_id:
            body["prompt_version_id"] = prompt_version_id
        return await self._request("POST", "/v1/workspace/favorites", json=body)

    async def update_favorite_entry(
        self, prompt_id: str, prompt_version_id: str
    ) -> dict:
        return await self._request(
            "PATCH",
            f"/v1/workspace/favorites/{prompt_id}",
            json={"prompt_version_id": prompt_version_id},
        )

    # ── Intents ──────────────────────────────────────────────────────────

    async def submit_intent(self, intent: dict) -> dict:
        return await self._request("POST", "/v1/intents", json=intent)

    async def list_intents(
        self,
        *,
        limit: int | None = None,
        offset: int | None = None,
        status: str | None = None,
    ) -> dict:
        params = {}
        if limit is not None:
            params["limit"] = str(limit)
        if offset is not None:
            params["offset"] = str(offset)
        if status:
            params["status"] = status
        return await self._request("GET", "/v1/intents", params=params)

    async def get_intent(self, intent_id: str) -> dict:
        return await self._request("GET", f"/v1/intents/{intent_id}")

    async def clarify_intent(
        self, intent_id: str, clarification_input: str
    ) -> dict:
        return await self._request(
            "POST",
            f"/v1/intents/{intent_id}/clarify",
            json={"clarification_input": clarification_input},
        )

    # ── Packages ─────────────────────────────────────────────────────────

    async def execute_package(
        self,
        package_id: str,
        *,
        input: str | None = None,
        model: str | None = None,
        provider_api_key: str | None = None,
    ) -> dict:
        body: dict[str, Any] = {}
        if input:
            body["input"] = input
        if model:
            body["model"] = model
        if provider_api_key:
            body["provider_api_key"] = provider_api_key
        return await self._request(
            "POST", f"/v1/packages/{package_id}/execute", json=body
        )

    # ── Notifications ────────────────────────────────────────────────────

    async def get_notifications(self, *, limit: int | None = None) -> dict:
        params = {}
        if limit is not None:
            params["limit"] = str(limit)
        return await self._request("GET", "/v1/notifications", params=params)

    async def mark_notifications_read(
        self, ids: list[str] | str
    ) -> dict:
        return await self._request(
            "POST", "/v1/notifications/read", json={"ids": ids}
        )

    # ── Analytics ────────────────────────────────────────────────────────

    async def get_creator_analytics(
        self,
        *,
        period_start: str | None = None,
        period_end: str | None = None,
    ) -> dict:
        params = {}
        if period_start:
            params["period_start"] = period_start
        if period_end:
            params["period_end"] = period_end
        return await self._request(
            "GET", "/v1/analytics/creator", params=params
        )

    # ── Usage ────────────────────────────────────────────────────────────

    async def get_usage_summary(self, *, days: int | None = None) -> dict:
        params = {}
        if days is not None:
            params["days"] = str(days)
        return await self._request("GET", "/v1/usage/summary", params=params)

    async def get_usage_by_day(self, *, days: int | None = None) -> dict:
        params = {}
        if days is not None:
            params["days"] = str(days)
        return await self._request("GET", "/v1/usage/by-day", params=params)

    # ── Feedback ─────────────────────────────────────────────────────────

    async def submit_feedback(self, feedback: dict) -> dict:
        return await self._request("POST", "/v1/feedback", json=feedback)

    # ── RPC ──────────────────────────────────────────────────────────────

    async def rpc(self, method: str, params: dict) -> Any:
        import time

        response = await self._request(
            "POST",
            "/v1/rpc",
            json={
                "jsonrpc": "2.0",
                "method": method,
                "params": params,
                "id": int(time.time() * 1000),
            },
        )
        if "error" in response:
            raise A21eError(
                response["error"].get("message", "RPC error"),
                code=response["error"].get("code"),
            )
        return response.get("result")
