"""Memories API namespace — the complex one with streaming, batch, file upload."""

from __future__ import annotations

import mimetypes
from pathlib import Path
from typing import Any, Literal, overload

import httpx

from goodmem._base import AsyncCRUDNamespace, CRUDNamespace
from goodmem._generated.models.batch_memory_creation_request import BatchMemoryCreationRequest
from goodmem._generated.models.batch_memory_deletion_request import BatchMemoryDeletionRequest
from goodmem._generated.models.batch_memory_response import BatchMemoryResponse
from goodmem._generated.models.batch_memory_retrieval_request import BatchMemoryRetrievalRequest
from goodmem._generated.models.memory import Memory
from goodmem._generated.models.memory_creation_request import MemoryCreationRequest
from goodmem._generated.models.memory_list_response import MemoryListResponse
from goodmem._generated.models.retrieve_memory_event import RetrieveMemoryEvent
from goodmem._generated.models.retrieve_memory_request import RetrieveMemoryRequest
from goodmem.errors import raise_for_status
from goodmem.pagination import AsyncPaginator, Paginator
from goodmem.streaming import AsyncRetrieveMemoryStream, RetrieveMemoryStream

_CHAT_POST_PROCESSOR_NAME = "com.goodmem.retrieval.postprocess.ChatPostProcessorFactory"

_POST_PROCESSOR_ARGS = frozenset({
    "llm_id", "reranker_id", "prompt", "sys_prompt", "gen_token_budget",
    "relevance_threshold", "llm_temp", "max_results", "chronological_resort",
})


def _build_retrieve_body(message: str, kwargs: dict[str, Any]) -> dict:
    """Build RetrieveMemoryRequest dict, handling convenience params."""
    # Convert space_ids shorthand → space_keys
    if "space_ids" in kwargs:
        space_ids = kwargs.pop("space_ids")
        if isinstance(space_ids, str):
            space_ids = [space_ids]
        kwargs.setdefault("space_keys", [{"space_id": sid} for sid in space_ids])

    # Flatten post-processor args → nested post_processor with generic name+config
    pp_kwargs = {k: kwargs.pop(k) for k in list(kwargs) if k in _POST_PROCESSOR_ARGS}
    if pp_kwargs:
        if "post_processor" in kwargs:
            raise ValueError("Cannot specify both 'post_processor' and individual post-processor args.")
        # Store raw post_processor dict — will be injected after model_validate
        raw_post_processor = {
            "name": _CHAT_POST_PROCESSOR_NAME,
            "config": pp_kwargs,
        }
    else:
        raw_post_processor = kwargs.pop("post_processor", None)

    kwargs["message"] = message
    body = RetrieveMemoryRequest.model_validate(kwargs)
    result = body.model_dump(by_alias=True, exclude_unset=True)

    # Inject post_processor directly into serialized body (bypasses strict typing
    # on config values, since the server accepts plain strings in the config map)
    if raw_post_processor is not None:
        result["postProcessor"] = raw_post_processor

    return result


def _infer_content_type(file_path: str | Path) -> str:
    """Infer MIME type from file extension."""
    mime_type, _ = mimetypes.guess_type(str(file_path))
    return mime_type or "application/octet-stream"


class MemoriesAPI(CRUDNamespace):
    _base_path = "/v1/memories"
    _create_model = MemoryCreationRequest
    _update_model = None
    _response_model = Memory
    _list_response_model = MemoryListResponse
    _list_items_field = "memories"

    def create(
        self,
        *,
        space_id: str | None = None,
        file_paths: str | list[str] | None = None,
        **kwargs: Any,
    ) -> Memory:
        """Create a memory. Supports text content or file upload.

        For text::

            client.memories.create(space_id="...", original_content="Hello", content_type="text/plain")

        For file upload::

            client.memories.create(space_id="...", file_paths="doc.pdf")
        """
        if space_id is not None:
            kwargs["space_id"] = space_id

        if file_paths is not None:
            return self._create_with_file(file_paths, kwargs)

        # Auto-infer content_type for text
        if "content_type" not in kwargs and "original_content" in kwargs:
            if isinstance(kwargs["original_content"], str):
                kwargs["content_type"] = "text/plain"

        body = self._create_model.model_validate(kwargs)
        resp = self._http.post(
            self._base_path,
            json=body.model_dump(by_alias=True, exclude_unset=True),
            headers=self._headers,
        )
        return self._handle_response(resp, self._response_model)

    def _create_with_file(self, file_paths: str | list[str], kwargs: dict) -> Memory:
        """Handle multipart file upload."""
        if isinstance(file_paths, str):
            file_paths = [file_paths]
        if len(file_paths) != 1:
            raise ValueError("Single file upload only. Use batch_create for multiple files.")

        path = Path(file_paths[0])
        if not path.exists():
            raise FileNotFoundError(f"File not found: {path}")

        kwargs.setdefault("content_type", _infer_content_type(path))

        request_data = self._create_model.model_validate(kwargs)
        with open(path, "rb") as f:
            resp = self._http.post(
                self._base_path,
                data={"request": request_data.model_dump_json(by_alias=True, exclude_unset=True)},
                files={"file": (path.name, f, kwargs["content_type"])},
                headers=self._headers,
            )
        return self._handle_response(resp, self._response_model)

    def list(self, space_id: str, **kwargs: Any) -> Paginator[Memory]:
        """List memories in a space."""
        return Paginator(
            fetch_fn=lambda token: self._http.get(
                f"/v1/spaces/{space_id}/memories",
                params={**kwargs, "nextToken": token} if token else kwargs or None,
                headers=self._headers,
            ),
            response_model=self._list_response_model,
            items_field=self._list_items_field,
        )

    def get_content(self, id: str) -> bytes:
        """Download raw memory content."""
        resp = self._http.get(
            f"{self._base_path}/{id}/content",
            headers=self._headers,
        )
        raise_for_status(resp)
        return resp.content

    @overload
    def retrieve(self, *, message: str, stream: Literal[False], **kwargs: Any) -> list[RetrieveMemoryEvent]: ...
    @overload
    def retrieve(self, *, message: str, stream: Literal[True] = ..., **kwargs: Any) -> RetrieveMemoryStream: ...

    def retrieve(
        self,
        *,
        message: str,
        stream: bool = True,
        **kwargs: Any,
    ) -> RetrieveMemoryStream | list[RetrieveMemoryEvent]:
        """Semantic memory retrieval.

        Args:
            message:      Query string for semantic search.
            stream:       If True (default), returns a context manager yielding events.
                          If False, collects all events and returns a list.
            space_ids:    Convenience — list of space ID strings (converted to space_keys).
            space_keys:   List of SpaceKey dicts with optional embedder weight overrides.
            llm_id:       Post-processor LLM ID (flattened convenience param).
            reranker_id:  Post-processor reranker ID (flattened convenience param).
            relevance_threshold, llm_temp, max_results, chronological_resort, prompt,
            sys_prompt, gen_token_budget: Additional post-processor params.
            context:      List of context items (text or binary).
            filter:       SQL-like filter expression.
        """
        body_dict = _build_retrieve_body(message, kwargs)
        if stream:
            return RetrieveMemoryStream(
                self._http,
                method="POST",
                url="/v1/memories:retrieve",
                json=body_dict,
                headers={**self._headers, "Accept": "application/x-ndjson"},
            )
        else:
            with RetrieveMemoryStream(
                self._http,
                method="POST",
                url="/v1/memories:retrieve",
                json=body_dict,
                headers={**self._headers, "Accept": "application/x-ndjson"},
            ) as s:
                return list(s)

    def batch_create(self, items: list[dict | MemoryCreationRequest]) -> BatchMemoryResponse:
        """Create multiple memories. Accepts a bare list of items."""
        # Auto-infer content_type for text items
        processed = []
        for item in items:
            if isinstance(item, dict):
                item = dict(item)  # copy
                if "content_type" not in item and "original_content" in item:
                    if isinstance(item["original_content"], str):
                        item["content_type"] = "text/plain"
                processed.append(item)
            else:
                processed.append(item)

        body = BatchMemoryCreationRequest.model_validate({"requests": processed})
        resp = self._http.post(
            f"{self._base_path}:batchCreate",
            json=body.model_dump(by_alias=True, exclude_unset=True),
            headers=self._headers,
        )
        return self._handle_response(resp, BatchMemoryResponse)

    def batch_get(self, ids: list[str]) -> BatchMemoryResponse:
        """Get multiple memories by ID. Accepts a bare list of IDs."""
        body = BatchMemoryRetrievalRequest.model_validate({"memory_ids": ids})
        resp = self._http.post(
            f"{self._base_path}:batchGet",
            json=body.model_dump(by_alias=True, exclude_unset=True),
            headers=self._headers,
        )
        return self._handle_response(resp, BatchMemoryResponse)

    def batch_delete(self, ids: list[str]) -> BatchMemoryResponse:
        """Delete multiple memories by ID. Accepts a bare list of IDs."""
        body = BatchMemoryDeletionRequest.model_validate({"memory_ids": ids})
        resp = self._http.post(
            f"{self._base_path}:batchDelete",
            json=body.model_dump(by_alias=True, exclude_unset=True),
            headers=self._headers,
        )
        return self._handle_response(resp, BatchMemoryResponse)


class AsyncMemoriesAPI(AsyncCRUDNamespace):
    _base_path = "/v1/memories"
    _create_model = MemoryCreationRequest
    _update_model = None
    _response_model = Memory
    _list_response_model = MemoryListResponse
    _list_items_field = "memories"

    async def create(
        self,
        *,
        space_id: str | None = None,
        file_paths: str | list[str] | None = None,
        **kwargs: Any,
    ) -> Memory:
        """Create a memory. Supports text content or file upload."""
        if space_id is not None:
            kwargs["space_id"] = space_id

        if file_paths is not None:
            return await self._create_with_file(file_paths, kwargs)

        if "content_type" not in kwargs and "original_content" in kwargs:
            if isinstance(kwargs["original_content"], str):
                kwargs["content_type"] = "text/plain"

        body = self._create_model.model_validate(kwargs)
        resp = await self._http.post(
            self._base_path,
            json=body.model_dump(by_alias=True, exclude_unset=True),
            headers=self._headers,
        )
        return self._handle_response(resp, self._response_model)

    async def _create_with_file(self, file_paths: str | list[str], kwargs: dict) -> Memory:
        if isinstance(file_paths, str):
            file_paths = [file_paths]
        if len(file_paths) != 1:
            raise ValueError("Single file upload only. Use batch_create for multiple files.")

        path = Path(file_paths[0])
        if not path.exists():
            raise FileNotFoundError(f"File not found: {path}")

        kwargs.setdefault("content_type", _infer_content_type(path))
        request_data = self._create_model.model_validate(kwargs)
        with open(path, "rb") as f:
            resp = await self._http.post(
                self._base_path,
                data={"request": request_data.model_dump_json(by_alias=True, exclude_unset=True)},
                files={"file": (path.name, f, kwargs["content_type"])},
                headers=self._headers,
            )
        return self._handle_response(resp, self._response_model)

    def list(self, space_id: str, **kwargs: Any) -> AsyncPaginator[Memory]:
        return AsyncPaginator(
            fetch_fn=lambda token: self._http.get(
                f"/v1/spaces/{space_id}/memories",
                params={**kwargs, "nextToken": token} if token else kwargs or None,
                headers=self._headers,
            ),
            response_model=self._list_response_model,
            items_field=self._list_items_field,
        )

    async def get_content(self, id: str) -> bytes:
        resp = await self._http.get(
            f"{self._base_path}/{id}/content",
            headers=self._headers,
        )
        raise_for_status(resp)
        return resp.content

    async def retrieve(
        self,
        *,
        message: str,
        stream: bool = True,
        **kwargs: Any,
    ) -> AsyncRetrieveMemoryStream | list[RetrieveMemoryEvent]:
        body_dict = _build_retrieve_body(message, kwargs)
        if stream:
            return AsyncRetrieveMemoryStream(
                self._http,
                method="POST",
                url="/v1/memories:retrieve",
                json=body_dict,
                headers={**self._headers, "Accept": "application/x-ndjson"},
            )
        else:
            async with AsyncRetrieveMemoryStream(
                self._http,
                method="POST",
                url="/v1/memories:retrieve",
                json=body_dict,
                headers={**self._headers, "Accept": "application/x-ndjson"},
            ) as s:
                return [event async for event in s]

    async def batch_create(self, items: list[dict | MemoryCreationRequest]) -> BatchMemoryResponse:
        processed = []
        for item in items:
            if isinstance(item, dict):
                item = dict(item)
                if "content_type" not in item and "original_content" in item:
                    if isinstance(item["original_content"], str):
                        item["content_type"] = "text/plain"
                processed.append(item)
            else:
                processed.append(item)
        body = BatchMemoryCreationRequest.model_validate({"requests": processed})
        resp = await self._http.post(
            f"{self._base_path}:batchCreate",
            json=body.model_dump(by_alias=True, exclude_unset=True),
            headers=self._headers,
        )
        return self._handle_response(resp, BatchMemoryResponse)

    async def batch_get(self, ids: list[str]) -> BatchMemoryResponse:
        body = BatchMemoryRetrievalRequest.model_validate({"memory_ids": ids})
        resp = await self._http.post(
            f"{self._base_path}:batchGet",
            json=body.model_dump(by_alias=True, exclude_unset=True),
            headers=self._headers,
        )
        return self._handle_response(resp, BatchMemoryResponse)

    async def batch_delete(self, ids: list[str]) -> BatchMemoryResponse:
        body = BatchMemoryDeletionRequest.model_validate({"memory_ids": ids})
        resp = await self._http.post(
            f"{self._base_path}:batchDelete",
            json=body.model_dump(by_alias=True, exclude_unset=True),
            headers=self._headers,
        )
        return self._handle_response(resp, BatchMemoryResponse)
