from __future__ import annotations

from importlib import import_module
from typing import Any, Dict, Optional, Protocol, Sequence, Union
from uuid import uuid4

from gaunt.enums import ScoringTag
from gaunt.models import Inputs, Outputs, TracePayload


class SyncHTTPClient(Protocol):
    def post(self, url: str, json: Dict[str, Any]) -> Any: ...

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


class AsyncHTTPClient(Protocol):
    async def post(self, url: str, json: Dict[str, Any]) -> Any: ...

    async def aclose(self) -> None: ...


class Client:
    """SDK client for sending trace logs to the Analytics Bank API."""

    def __init__(
        self,
        api_key: str,
        base_url: str,
        timeout: float = 10.0,
        sync_http_client: Optional[SyncHTTPClient] = None,
        async_http_client: Optional[AsyncHTTPClient] = None,
    ) -> None:
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout

        if sync_http_client is not None:
            self._sync_http_client = sync_http_client
        else:
            httpx = import_module("httpx")
            self._sync_http_client = httpx.Client(
                base_url=self.base_url,
                timeout=self.timeout,
                headers=self._default_headers,
            )

        if async_http_client is not None:
            self._async_http_client = async_http_client
        else:
            httpx = import_module("httpx")
            self._async_http_client = httpx.AsyncClient(
                base_url=self.base_url,
                timeout=self.timeout,
                headers=self._default_headers,
            )

    @property
    def _default_headers(self) -> Dict[str, str]:
        return {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

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

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

    def _normalize_tags(
        self, scoring_tags: Optional[Sequence[Union[ScoringTag, str]]]
    ) -> list[str]:
        if not scoring_tags:
            return []
        return [tag.value if isinstance(tag, ScoringTag) else tag for tag in scoring_tags]

    def _build_payload(
        self,
        project_id: str,
        inputs: Union[Inputs, Dict[str, Any]],
        outputs: Union[Outputs, Dict[str, Any]],
        expected: Optional[Any] = None,
        metadata: Optional[Dict[str, Any]] = None,
        scoring_tags: Optional[Sequence[Union[ScoringTag, str]]] = None,
        trace_id: Optional[str] = None,
    ) -> TracePayload:
        payload = TracePayload(
            project_id=project_id,
            trace_id=trace_id or str(uuid4()),
            inputs=inputs,
            outputs=outputs,
            expected=expected,
            metadata=metadata,
            config={"scoring_tags": self._normalize_tags(scoring_tags)},
        )
        return payload

    def log(
        self,
        project_id: str,
        inputs: Union[Inputs, Dict[str, Any]],
        outputs: Union[Outputs, Dict[str, Any]],
        expected: Optional[Any] = None,
        metadata: Optional[Dict[str, Any]] = None,
        scoring_tags: Optional[Sequence[Union[ScoringTag, str]]] = None,
        trace_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        payload = self._build_payload(
            project_id=project_id,
            trace_id=trace_id,
            inputs=inputs,
            outputs=outputs,
            expected=expected,
            metadata=metadata,
            scoring_tags=scoring_tags,
        )
        response = self._sync_http_client.post(
            "/logs", json=payload.model_dump(by_alias=True, exclude_none=True)
        )
        response.raise_for_status()
        if response.content:
            return response.json()
        return {"status": "ok"}

    async def alog(
        self,
        project_id: str,
        inputs: Union[Inputs, Dict[str, Any]],
        outputs: Union[Outputs, Dict[str, Any]],
        expected: Optional[Any] = None,
        metadata: Optional[Dict[str, Any]] = None,
        scoring_tags: Optional[Sequence[Union[ScoringTag, str]]] = None,
        trace_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        payload = self._build_payload(
            project_id=project_id,
            trace_id=trace_id,
            inputs=inputs,
            outputs=outputs,
            expected=expected,
            metadata=metadata,
            scoring_tags=scoring_tags,
        )
        response = await self._async_http_client.post(
            "/logs", json=payload.model_dump(by_alias=True, exclude_none=True)
        )
        response.raise_for_status()
        if response.content:
            return response.json()
        return {"status": "ok"}
