"""
Cognee Client abstraction that supports both direct function calls and HTTP API calls.

This module provides a unified interface for interacting with Cognee, supporting:
- Direct mode: Directly imports and calls cognee functions (default behavior)
- API mode: Makes HTTP requests to a running Cognee FastAPI server
"""

import sys
import hashlib
from typing import Optional, Any, List, Dict
from uuid import UUID
from contextlib import redirect_stdout
import httpx
from cognee.shared.logging_utils import get_logger
import json

try:
    from .server_utils import normalize_delete_mode
except ImportError:
    from server_utils import normalize_delete_mode

try:
    from .retrieval_utils import get_chunk_neighbors_from_graph, get_document_from_graph
except ImportError:
    from retrieval_utils import get_chunk_neighbors_from_graph, get_document_from_graph

logger = get_logger()


class CogneeClient:
    """
    Unified client for interacting with Cognee via direct calls or HTTP API.

    Parameters
    ----------
    api_url : str, optional
        Base URL of the Cognee API server (e.g., "http://localhost:8000").
        If None, uses direct cognee function calls.
    api_token : str, optional
        Authentication token for the API (optional, required if API has authentication enabled).
    """

    def __init__(self, api_url: Optional[str] = None, api_token: Optional[str] = None):
        self.api_url = api_url.rstrip("/") if api_url else None
        self.api_token = api_token
        self.use_api = bool(api_url)

        # Extract tenant ID from tenant URL pattern: tenant-<uuid>.*.cognee.ai
        self.tenant_id: Optional[str] = None
        if self.api_url:
            import re

            match = re.search(r"tenant-([0-9a-f-]{36})", self.api_url)
            if match:
                self.tenant_id = match.group(1)

        if self.use_api:
            logger.info(f"Cognee client initialized in API mode: {self.api_url}")
            if self.tenant_id:
                logger.info(f"Tenant ID extracted from URL: {self.tenant_id}")
            self.client = httpx.AsyncClient(timeout=300.0)  # 5 minute timeout for long operations
        else:
            logger.info("Cognee client initialized in direct mode")
            # Import cognee only if we're using direct mode
            import cognee as _cognee

            self.cognee = _cognee

    def _get_headers(self, include_content_type: bool = True) -> Dict[str, str]:
        """Get headers for API requests.

        Uses X-Api-Key + X-Tenant-Id for tenant APIs (cloud),
        falls back to Bearer token for local/self-hosted backends.
        """
        headers: Dict[str, str] = {}
        if include_content_type:
            headers["Content-Type"] = "application/json"
        if self.api_token:
            if self.tenant_id:
                headers["X-Api-Key"] = self.api_token
                headers["X-Tenant-Id"] = self.tenant_id
            else:
                headers["Authorization"] = f"Bearer {self.api_token}"
        return headers

    @staticmethod
    def _json_or_success(response: httpx.Response) -> Dict[str, Any]:
        """Return a JSON body when present, otherwise a generic success shape."""
        if not response.content:
            return {"status": "success"}
        try:
            parsed = response.json()
        except ValueError:
            return {"status": "success", "message": response.text}
        if isinstance(parsed, dict):
            return parsed
        return {"status": "success", "result": parsed}

    @staticmethod
    def _text_upload(data: Any) -> Dict[str, tuple[str, str, str]]:
        """Create a content-addressed text upload for API-mode ingestion."""
        content = str(data)
        digest = hashlib.md5(content.encode("utf-8")).hexdigest()
        return {"data": (f"text_{digest}.txt", content, "text/plain")}

    async def add(
        self, data: Any, dataset_name: str = "main_dataset", node_set: Optional[List[str]] = None
    ) -> Dict[str, Any]:
        """
        Add data to Cognee for processing.

        Parameters
        ----------
        data : Any
            Data to add (text, file path, etc.)
        dataset_name : str
            Name of the dataset to add data to
        node_set : List[str], optional
            List of node identifiers for graph organization

        Returns
        -------
        Dict[str, Any]
            Result of the add operation
        """
        if self.use_api:
            endpoint = f"{self.api_url}/api/v1/add"

            files = self._text_upload(data)
            form_data = {
                "datasetName": dataset_name,
            }
            if node_set is not None:
                form_data["node_set"] = json.dumps(node_set)

            response = await self.client.post(
                endpoint,
                files=files,
                data=form_data,
                headers=self._get_headers(include_content_type=False),
            )
            response.raise_for_status()
            return response.json()
        else:
            with redirect_stdout(sys.stderr):
                await self.cognee.add(data, dataset_name=dataset_name, node_set=node_set)
                return {"status": "success", "message": "Data added successfully"}

    async def cognify(
        self,
        datasets: Optional[List[str]] = None,
        custom_prompt: Optional[str] = None,
        graph_model: Any = None,
    ) -> Dict[str, Any]:
        """
        Transform data into a knowledge graph.

        Parameters
        ----------
        datasets : List[str], optional
            List of dataset names to process
        custom_prompt : str, optional
            Custom prompt for entity extraction
        graph_model : Any, optional
            Custom graph model (only used in direct mode)

        Returns
        -------
        Dict[str, Any]
            Result of the cognify operation
        """
        if self.use_api:
            # API mode: Make HTTP request
            endpoint = f"{self.api_url}/api/v1/cognify"
            payload = {
                "datasets": datasets or ["main_dataset"],
                "run_in_background": False,
            }
            if custom_prompt:
                payload["custom_prompt"] = custom_prompt

            response = await self.client.post(endpoint, json=payload, headers=self._get_headers())
            response.raise_for_status()
            return response.json()
        else:
            # Direct mode: Call cognee directly
            with redirect_stdout(sys.stderr):
                kwargs = {}
                if datasets:
                    kwargs["datasets"] = datasets
                if custom_prompt:
                    kwargs["custom_prompt"] = custom_prompt
                if graph_model:
                    kwargs["graph_model"] = graph_model

                await self.cognee.cognify(**kwargs)
                return {"status": "success", "message": "Cognify completed successfully"}

    async def search(
        self,
        query_text: str,
        query_type: str,
        datasets: Optional[List[str]] = None,
        system_prompt: Optional[str] = None,
        top_k: int = 15,
    ) -> Any:
        """
        Search the knowledge graph.

        Parameters
        ----------
        query_text : str
            The search query
        query_type : str
            Type of search (e.g., "GRAPH_COMPLETION", "INSIGHTS", etc.)
        datasets : List[str], optional
            List of datasets to search
        system_prompt : str, optional
            System prompt for completion searches
        top_k : int
            Maximum number of results

        Returns
        -------
        Any
            Search results
        """
        if self.use_api:
            # API mode: Make HTTP request
            endpoint = f"{self.api_url}/api/v1/search"
            payload = {"query": query_text, "search_type": query_type.upper(), "top_k": top_k}
            if datasets:
                payload["datasets"] = datasets
            if system_prompt:
                payload["system_prompt"] = system_prompt

            response = await self.client.post(endpoint, json=payload, headers=self._get_headers())
            response.raise_for_status()
            return response.json()
        else:
            # Direct mode: Call cognee directly
            from cognee.modules.search.types import SearchType

            with redirect_stdout(sys.stderr):
                search_kwargs = {
                    "query_type": SearchType[query_type.upper()],
                    "query_text": query_text,
                    "top_k": top_k,
                }
                if datasets:
                    search_kwargs["datasets"] = datasets
                if system_prompt:
                    search_kwargs["system_prompt"] = system_prompt
                results = await self.cognee.search(**search_kwargs)
                return results

    async def delete(self, data_id: UUID, dataset_id: UUID, mode: str = "soft") -> Dict[str, Any]:
        """
        Delete data from a dataset.

        Parameters
        ----------
        data_id : UUID
            ID of the data to delete
        dataset_id : UUID
            ID of the dataset containing the data

        Returns
        -------
        Dict[str, Any]
            Result of the deletion
        """
        normalized_mode = normalize_delete_mode(mode)

        if self.use_api:
            # The deprecated delete endpoint still carries the mode contract.
            # Fall back to the datasets endpoint for older backends that removed it.
            endpoint = f"{self.api_url}/api/v1/delete"
            response = await self.client.delete(
                endpoint,
                params={
                    "data_id": str(data_id),
                    "dataset_id": str(dataset_id),
                    "mode": normalized_mode,
                },
                headers=self._get_headers(),
            )
            if response.status_code in {404, 405}:
                endpoint = f"{self.api_url}/api/v1/datasets/{str(dataset_id)}/data/{str(data_id)}"
                response = await self.client.delete(endpoint, headers=self._get_headers())
            response.raise_for_status()
            return self._json_or_success(response)
        else:
            # Direct mode: Call cognee directly
            from cognee.modules.users.methods import get_default_user

            with redirect_stdout(sys.stderr):
                user = await get_default_user()
                result = await self.cognee.datasets.delete_data(
                    dataset_id=dataset_id,
                    data_id=data_id,
                    mode=normalized_mode,
                    user=user,
                )
                return result or {"status": "success"}

    async def prune_data(self) -> Dict[str, Any]:
        """
        Prune all data from the knowledge graph.

        Returns
        -------
        Dict[str, Any]
            Result of the prune operation
        """
        if self.use_api:
            # Note: The API doesn't expose a prune endpoint, so we'll need to handle this
            # For now, raise an error
            raise NotImplementedError("Prune operation is not available via API")
        else:
            # Direct mode: Call cognee directly
            with redirect_stdout(sys.stderr):
                await self.cognee.prune.prune_data()
                return {"status": "success", "message": "Data pruned successfully"}

    async def prune_system(self, metadata: bool = True) -> Dict[str, Any]:
        """
        Prune system data from the knowledge graph.

        Parameters
        ----------
        metadata : bool
            Whether to prune metadata

        Returns
        -------
        Dict[str, Any]
            Result of the prune operation
        """
        if self.use_api:
            # Note: The API doesn't expose a prune endpoint
            raise NotImplementedError("Prune system operation is not available via API")
        else:
            # Direct mode: Call cognee directly
            with redirect_stdout(sys.stderr):
                await self.cognee.prune.prune_system(metadata=metadata)
                return {"status": "success", "message": "System pruned successfully"}

    async def get_pipeline_status(
        self, dataset_ids: List[UUID], pipeline_name: str
    ) -> Dict[str, Any]:
        """
        Get the status of a pipeline run.

        Parameters
        ----------
        dataset_ids : List[UUID]
            List of dataset IDs
        pipeline_name : str
            Name of the pipeline

        Returns
        -------
        Dict[str, Any]
            Status information keyed by dataset ID
        """
        if self.use_api:
            # API mode: query the server's dataset-status endpoint, which
            # reports the pipeline run state keyed by dataset id.
            endpoint = f"{self.api_url}/api/v1/datasets/status"
            params = [("dataset", str(d)) for d in dataset_ids]
            response = await self.client.get(endpoint, params=params, headers=self._get_headers())
            response.raise_for_status()
            return response.json()
        else:
            # Direct mode: Call cognee directly
            from cognee.modules.pipelines.operations.get_pipeline_status import get_pipeline_status

            with redirect_stdout(sys.stderr):
                status = await get_pipeline_status(dataset_ids, pipeline_name)
                return status

    async def list_datasets(self) -> List[Dict[str, Any]]:
        """
        List all datasets.

        Returns
        -------
        List[Dict[str, Any]]
            List of datasets
        """
        if self.use_api:
            # API mode: Make HTTP request
            endpoint = f"{self.api_url}/api/v1/datasets"
            response = await self.client.get(endpoint, headers=self._get_headers())
            response.raise_for_status()
            return response.json()
        else:
            # Direct mode: Call cognee directly
            from cognee.modules.users.methods import get_default_user
            from cognee.modules.data.methods import get_datasets

            with redirect_stdout(sys.stderr):
                user = await get_default_user()
                datasets = await get_datasets(user.id)
                return [
                    {"id": str(d.id), "name": d.name, "created_at": str(d.created_at)}
                    for d in datasets
                ]

    async def get_document(
        self,
        document_id: str,
        include_metadata: bool = True,
        max_chunks: int = 0,
    ) -> Dict[str, Any]:
        """Retrieve a full document with its chunks from the graph database."""
        if self.use_api:
            raise NotImplementedError("get_document is not available in API mode")

        from cognee.infrastructure.databases.unified import get_unified_engine

        with redirect_stdout(sys.stderr):
            unified = await get_unified_engine()
            return await get_document_from_graph(
                unified.graph,
                document_id,
                include_metadata=include_metadata,
                max_chunks=max_chunks,
            )

    async def get_chunk_neighbors(
        self,
        chunk_id: str,
        neighbor_count: int = 2,
        include_target: bool = True,
        direction: str = "both",
    ) -> Dict[str, Any]:
        """Retrieve neighboring chunks around a target chunk from its parent document."""
        if self.use_api:
            raise NotImplementedError("get_chunk_neighbors is not available in API mode")

        from cognee.infrastructure.databases.unified import get_unified_engine

        with redirect_stdout(sys.stderr):
            unified = await get_unified_engine()
            return await get_chunk_neighbors_from_graph(
                unified.graph,
                chunk_id,
                neighbor_count=neighbor_count,
                include_target=include_target,
                direction=direction,
            )

    # -- V2 API methods -----------------------------------------------------

    async def remember(
        self,
        data: Any,
        dataset_name: str = "main_dataset",
        session_id: Optional[str] = None,
        custom_prompt: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Store data in memory via remember().

        With session_id: stores in session cache only (fast).
        Without session_id: full add + cognify pipeline (permanent).
        """
        if self.use_api:
            if session_id:
                if custom_prompt:
                    logger.warning(
                        "remember: custom_prompt is not supported with session_id in API mode "
                        "(the /remember/entry endpoint does not forward custom_prompt)"
                    )
                    raise ValueError(
                        "custom_prompt is not supported when session_id is provided in API mode"
                    )
                # Session mode: POST a JSON QAEntry so the backend receives
                # real text, not a multipart-file placeholder that triggers
                # the _SESSION_PLACEHOLDER_PREFIXES skip in _add_to_session.
                endpoint = f"{self.api_url}/api/v1/remember/entry"
                payload = {
                    "entry": {
                        "type": "qa",
                        "question": "",
                        "answer": str(data),
                        "context": "",
                    },
                    "dataset_name": dataset_name,
                    "session_id": session_id,
                }
                response = await self.client.post(
                    endpoint,
                    json=payload,
                    headers=self._get_headers(),
                )
                response.raise_for_status()
                return response.json()

            endpoint = f"{self.api_url}/api/v1/remember"
            files = self._text_upload(data)
            form_data = {"datasetName": dataset_name}
            if custom_prompt:
                form_data["custom_prompt"] = custom_prompt
            response = await self.client.post(
                endpoint,
                files=files,
                data=form_data,
                headers=self._get_headers(include_content_type=False),
            )
            response.raise_for_status()
            return response.json()
        else:
            with redirect_stdout(sys.stderr):
                kwargs = {
                    "data": data,
                    "dataset_name": dataset_name,
                }
                if session_id:
                    kwargs["session_id"] = session_id
                if custom_prompt:
                    kwargs["custom_prompt"] = custom_prompt
                result = await self.cognee.remember(**kwargs)
                return {
                    "status": getattr(result, "status", "completed"),
                    "dataset_name": dataset_name,
                    "session_id": session_id,
                }

    async def recall(
        self,
        query_text: str,
        search_type: Optional[str] = None,
        datasets: Optional[List[str]] = None,
        session_id: Optional[str] = None,
        top_k: int = 15,
    ) -> Any:
        """Search memory via recall() with auto-routing and session awareness."""
        if self.use_api:
            endpoint = f"{self.api_url}/api/v1/recall"
            payload = {"query": query_text, "top_k": top_k, "search_type": None}
            if search_type:
                payload["search_type"] = search_type.upper()
            if datasets:
                payload["datasets"] = datasets
            if session_id:
                payload["session_id"] = session_id
            response = await self.client.post(endpoint, json=payload, headers=self._get_headers())
            response.raise_for_status()
            return response.json()
        else:
            with redirect_stdout(sys.stderr):
                kwargs = {"top_k": top_k, "auto_route": True}
                if search_type:
                    from cognee.modules.search.types import SearchType

                    kwargs["query_type"] = SearchType[search_type.upper()]
                if datasets:
                    kwargs["datasets"] = datasets
                if session_id:
                    kwargs["session_id"] = session_id
                return await self.cognee.recall(query_text=query_text, **kwargs)

    async def forget(
        self,
        dataset: Optional[str] = None,
        data_id: Optional[UUID] = None,
        dataset_id: Optional[UUID] = None,
        everything: bool = False,
        memory_only: bool = False,
    ) -> Dict[str, Any]:
        """Delete data via forget().

        Bug fix (kb#70): this method previously dropped `data_id`,
        `dataset_id`, and `memory_only` on the floor, so entry-level
        deletion was impossible through the MCP surface even though the
        cognee API's /api/v1/forget endpoint has always supported it
        (dataset/datasetId + dataId). Forward all fields it accepts.
        """
        if self.use_api:
            endpoint = f"{self.api_url}/api/v1/forget"
            payload = {"everything": everything, "memory_only": memory_only}
            if dataset:
                payload["dataset"] = dataset
            if dataset_id:
                payload["dataset_id"] = str(dataset_id)
            if data_id:
                payload["data_id"] = str(data_id)
            response = await self.client.post(endpoint, json=payload, headers=self._get_headers())
            response.raise_for_status()
            return response.json()
        else:
            with redirect_stdout(sys.stderr):
                return await self.cognee.forget(
                    dataset=dataset,
                    dataset_id=dataset_id,
                    data_id=data_id,
                    everything=everything,
                    memory_only=memory_only,
                )

    async def improve(
        self,
        dataset_name: str = "main_dataset",
        session_ids: Optional[List[str]] = None,
    ) -> Dict[str, Any]:
        """Enrich knowledge graph and bridge session data via improve()."""
        if self.use_api:
            endpoint = f"{self.api_url}/api/v1/improve"
            payload = {"dataset_name": dataset_name}
            if session_ids:
                payload["session_ids"] = session_ids
            response = await self.client.post(endpoint, json=payload, headers=self._get_headers())
            response.raise_for_status()
            return response.json()
        else:
            with redirect_stdout(sys.stderr):
                kwargs = {"dataset": dataset_name}
                if session_ids:
                    kwargs["session_ids"] = session_ids
                result = await self.cognee.improve(**kwargs)
                return {"status": "success", "result": str(result)}

    async def close(self):
        """Close the HTTP client if in API mode."""
        if self.use_api and hasattr(self, "client"):
            await self.client.aclose()
