from collections.abc import AsyncIterator
from contextlib import asynccontextmanager

import aiohttp
import orjson
from aiobaseclient import BaseStandardClient
from aiobaseclient.exceptions import NotFoundError
from izihawa_loglib.request_context import RequestContext
from izihawa_utils.common import filter_none

from spacefrontiers.clients.types import (
    ConversationRequest,
    ConversationResponse,
    SearchResponse,
)


def _detect_source(doc: dict) -> str:
    """Detect the document source from its type field."""
    doc_type = doc.get("type", "")
    if doc_type in ("submission", "comment", "message"):
        return "social"
    return "documents"


class SearchApiClient(BaseStandardClient):
    """Client for interacting with the Search API.

    A client that handles communication with the Search API endpoints, including authentication
    and request handling.

    Args:
        base_url (str): The base URL of the Search API.
        api_key (str | None, optional): API key for authentication. Defaults to None.
        user_id (str | int | None, optional): User ID of the user. Defaults to None.
        max_retries (int, optional): Maximum number of retry attempts. Defaults to 2.
        retry_delay (float, optional): Delay between retry attempts in seconds. Defaults to 0.5.
    """

    def __init__(
        self,
        base_url: str = "https://api.spacefrontiers.org",
        api_key: str | None = None,
        user_id: str | int | None = None,
        max_retries: int = 2,
        retry_delay: float = 0.5,
        default_headers: dict[str, str] | None = None,
    ):
        """Initialize the SearchApiClient.

        Args:
            base_url (str): The base URL of the Search API.
            api_key (str | None, optional): API key for authentication. Defaults to None.
            user_id (str | int | None, optional): User ID of the user. Defaults to None.
            max_retries (int, optional): Maximum number of retry attempts. Defaults to 2.
            retry_delay (float, optional): Delay between retry attempts in seconds. Defaults to 0.5.
            default_headers (dict[str, str] | None, optional): Default headers for requests. Defaults to None.
        """
        if default_headers is None:
            default_headers = {}
        default_headers.update(self._build_headers(api_key, user_id))
        super().__init__(
            base_url=base_url,
            default_headers=default_headers,
            max_retries=max_retries,
            retry_delay=retry_delay,
        )

    def _build_headers(
        self,
        api_key: str | None = None,
        user_id: str | int | None = None,
    ) -> dict[str, str]:
        headers = {}
        if api_key is not None:
            headers["X-Api-Key"] = api_key
        if user_id is not None:
            headers["X-User-Id"] = str(user_id)
        return headers

    async def post_conversations(
        self,
        conversation_request: ConversationRequest,
        timeout: float = 60.0,
        request_context: RequestContext | None = None,
        api_key: str | None = None,
        user_id: str | int | None = None,
    ) -> ConversationResponse:
        """Execute a conversation search request.

        Args:
            conversation_request (ConversationRequest): The conversation request parameters.
            timeout (float, optional): Request timeout in seconds. Defaults to 60.0.
            request_context (RequestContext | None, optional): Context for the request. Defaults to None.

        Returns:
            ConversationResponse: The conversation response
        """
        response = await self.post(
            "/v2/conversations/",
            json=filter_none(conversation_request.model_dump(exclude_none=True)),
            request_context=request_context,
            timeout=timeout,
            headers=self._build_headers(api_key, user_id),
        )
        return ConversationResponse(**response)

    @asynccontextmanager
    async def post_conversations_stream(
        self,
        conversation_request: ConversationRequest,
        timeout: float = 300.0,
        request_context: RequestContext | None = None,
        api_key: str | None = None,
        user_id: str | int | None = None,
    ) -> AsyncIterator[AsyncIterator[dict]]:
        """Stream a conversation request, yielding parsed SSE events.

        Usage::

            async with client.post_conversations_stream(req) as events:
                async for event in events:
                    if event["type"] == "delta":
                        print(event["text"])
                    elif event["type"] == "response":
                        final = event["response"]
        """
        url = f"{self.base_url}/v2/conversations/stream/"
        headers = {
            "Content-Type": "application/json",
            **self.default_headers,
            **self._build_headers(api_key, user_id),
        }
        if request_context is not None:
            for param in self.request_context_bypass or ():
                value = request_context.default_fields.get(param)
                if value is not None:
                    headers[f"x-{param.replace('_', '-')}"] = value
        body = orjson.dumps(
            filter_none(conversation_request.model_dump(exclude_none=True))
        )
        client_timeout = aiohttp.ClientTimeout(total=timeout)
        async with aiohttp.ClientSession(
            connector=self.connector,
            connector_owner=False,
            timeout=client_timeout,
        ) as session:
            async with session.post(url, data=body, headers=headers) as resp:
                if resp.status != 200:
                    text = await resp.text()
                    raise Exception(
                        f"Conversation stream failed ({resp.status}): {text}"
                    )

                async def _iter_events() -> AsyncIterator[dict]:
                    buf = b""
                    async for chunk in resp.content.iter_any():
                        buf += chunk
                        while b"\n" in buf:
                            line, buf = buf.split(b"\n", 1)
                            line = line.strip()
                            if not line:
                                continue
                            try:
                                yield orjson.loads(line)
                            except orjson.JSONDecodeError:
                                continue
                    if buf.strip():
                        try:
                            yield orjson.loads(buf.strip())
                        except orjson.JSONDecodeError:
                            pass

                yield _iter_events()

    async def search_v2(
        self,
        query: str,
        limit: int = 10,
        offset: int = 0,
        mode: str = "sparse",
        filter_types: list[str] | None = None,
        filter_languages: list[str] | None = None,
        filter_publisher: str | None = None,
        filter_issns: list[str] | None = None,
        deduplicate: bool = True,
        request_context: RequestContext | None = None,
    ) -> SearchResponse:
        """Search using the v2 Hermes-based search endpoint.

        Supports special query prefixes like ``orcid:<ID>``, ``issn:<ISSN>``,
        ``rd:<DOI>`` which are handled server-side.
        """
        body: dict = {
            "query": query,
            "limit": limit,
            "offset": offset,
            "mode": mode,
            "deduplicate": deduplicate,
            "combiner": "weighted_top_k",
            "combiner_top_k": 3,
            "combiner_decay": 0.7,
            "heap_factor": 0.7,
            "pruning": 0.8,
        }
        if filter_types:
            body["filter_types"] = filter_types
        if filter_languages:
            body["filter_languages"] = filter_languages
        if filter_publisher:
            body["filter_publisher"] = filter_publisher
        if filter_issns:
            body["filter_issns"] = filter_issns
        response = await self.post(
            "/v2/search/",
            json=body,
            request_context=request_context,
        )
        hits = response.get("hits") or []
        search_documents = []
        for hit in hits:
            doc = hit.get("document", {})
            doc["id"] = hit.get("id", "")
            snippets = [
                {
                    "field": s.get("field", "content"),
                    "text": s.get("text", ""),
                    "score": s.get("score", 0.0),
                }
                for s in (hit.get("snippets") or [])
            ]
            search_documents.append(
                {
                    "source": _detect_source(doc),
                    "document": doc,
                    "snippets": snippets,
                    "score": hit.get("score", 0.0),
                }
            )
        return SearchResponse(
            search_documents=search_documents,
            count=len(search_documents),
            has_next=response.get("has_next", False),
            total_count=response.get("total_hits"),
        )

    async def query_documents(
        self,
        metadata_filters: dict[str, list[str]],
        filter_types: list[str] | None = None,
        limit: int = 10,
        offset: int = 0,
        request_context: RequestContext | None = None,
    ) -> SearchResponse:
        """Query documents by metadata filters (e.g. ORCID) via AlloyDB."""
        body: dict = {
            "metadata_filters": metadata_filters,
            "limit": limit,
            "offset": offset,
        }
        if filter_types:
            body["filter_types"] = filter_types
        response = await self.post(
            "/internal/v2/search/documents/query",
            json=body,
            request_context=request_context,
        )
        hits = response.get("hits") or []
        search_documents = []
        for hit in hits:
            doc = hit.get("document", {})
            doc["id"] = hit.get("id", "")
            search_documents.append(
                {
                    "source": _detect_source(doc),
                    "document": doc,
                    "snippets": [],
                    "score": hit.get("score", 0.0),
                }
            )
        return SearchResponse(
            search_documents=search_documents,
            count=len(search_documents),
            has_next=response.get("has_next", False),
            total_count=response.get("total_hits"),
        )

    async def similar_v2(
        self,
        document_id: str,
        limit: int = 10,
        filter_types: list[str] | None = None,
        filter_languages: list[str] | None = None,
        request_context: RequestContext | None = None,
        api_key: str | None = None,
        user_id: str | int | None = None,
    ) -> SearchResponse:
        body: dict = {"document_id": document_id, "limit": limit}
        if filter_types:
            body["filter_types"] = filter_types
        if filter_languages:
            body["filter_languages"] = filter_languages
        response = await self.post(
            "/v2/search/similar",
            json=body,
            request_context=request_context,
            headers=self._build_headers(api_key, user_id),
        )
        # Convert V2SearchResponse → SearchResponse
        hits = response.get("hits") or []
        search_documents = []
        for hit in hits:
            doc = hit.get("document", {})
            doc["id"] = hit.get("id", "")
            snippets = [
                {
                    "field": s.get("field", "content"),
                    "text": s.get("text", ""),
                    "score": s.get("score", 0.0),
                }
                for s in (hit.get("snippets") or [])
            ]
            search_documents.append(
                {
                    "source": _detect_source(doc),
                    "document": doc,
                    "snippets": snippets,
                    "score": hit.get("score", 0.0),
                }
            )
        return SearchResponse(
            search_documents=search_documents,
            count=len(search_documents),
            has_next=response.get("has_next", False),
            total_count=response.get("total_hits"),
        )

    async def get_document_v2(
        self,
        document_id: str,
        request_context: RequestContext | None = None,
        timeout: float = 60.0,
    ) -> dict | None:
        """Fetch a document by ID from alloydb-documents (v2).

        Returns the document in v2 format: doc["id"] is a flat string,
        doc["uris"] is a list of URI strings.
        Returns None if not found.
        """
        try:
            response = await self.get(
                f"/v2/search/documents/{document_id}",
                request_context=request_context,
                timeout=timeout,
            )
        except NotFoundError:
            return None
        doc = dict(response.get("document", {}))
        doc["id"] = response.get("id", "")
        doc["uris"] = response.get("uris", [])
        return doc

    async def fetch_by_uris(
        self,
        uris: list[str],
        limit: int = 100,
        request_context: RequestContext | None = None,
        timeout: float = 60.0,
    ) -> SearchResponse:
        """Fetch multiple documents by their URIs from alloydb-documents (v2).

        Args:
            uris: List of document URIs (e.g. ["doi://10.xxx", "isbn://..."]).
            limit: Maximum number of documents to return.
            request_context: Context for the request with tracking information.
            timeout: Request timeout in seconds.

        Returns:
            SearchResponse with matched documents.
        """
        response = await self.post(
            "/internal/v2/search/documents/by-uris",
            json={"uris": uris, "limit": limit},
            timeout=timeout,
            request_context=request_context,
        )
        hits = response.get("hits") or []
        search_documents = []
        for hit in hits:
            doc = hit.get("document", {})
            doc["id"] = hit.get("id", "")
            snippets = [
                {
                    "field": s.get("field", "content"),
                    "text": s.get("text", ""),
                    "score": s.get("score", 0.0),
                }
                for s in (hit.get("snippets") or [])
            ]
            search_documents.append(
                {
                    "source": _detect_source(doc),
                    "document": doc,
                    "snippets": snippets,
                    "score": hit.get("score", 0.0),
                }
            )
        return SearchResponse(
            search_documents=search_documents,
            count=len(search_documents),
            has_next=response.get("has_next", False),
            total_count=response.get("total_hits"),
        )
