"""
HieromemCallbackHandler - LangChain callback for auto-persisting conversations.

Automatically saves conversation turns to HIEROMEM.
"""
from typing import Any, Dict, List, Optional
from uuid import UUID
import httpx

try:
    from langchain.callbacks.base import BaseCallbackHandler
    from langchain.schema import LLMResult
    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False
    BaseCallbackHandler = object
    LLMResult = object


class HieromemCallbackHandler(BaseCallbackHandler):  # type: ignore[misc]
    """
    LangChain callback handler that auto-persists to HIEROMEM.

    Usage:
        from hieromem.integrations.langchain import HieromemCallbackHandler

        handler = HieromemCallbackHandler(
            hieromem_url="http://localhost:8000",
            namespace="conversations"
        )

        chain = LLMChain(llm=llm, callbacks=[handler])
    """

    def __init__(
        self,
        hieromem_url: str = "http://localhost:8000",
        api_key: Optional[str] = None,
        namespace: str = "default",
        persist_inputs: bool = True,
        persist_outputs: bool = True,
    ):
        if not LANGCHAIN_AVAILABLE:
            raise ImportError(
                "LangChain is required for HieromemCallbackHandler. "
                "Install with: pip install langchain"
            )
        self.hieromem_url = hieromem_url
        self.api_key = api_key
        self.namespace = namespace
        self.persist_inputs = persist_inputs
        self.persist_outputs = persist_outputs
        self._client = httpx.Client(
            base_url=hieromem_url,
            headers={"Authorization": f"Bearer {api_key}"} if api_key else {},
            timeout=30.0,
        )

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Persist prompts/inputs to HIEROMEM."""
        if not self.persist_inputs:
            return

        for prompt in prompts:
            self._insert_episode(
                content=prompt,
                meta={
                    "type": "llm_input",
                    "run_id": str(run_id),
                },
            )

    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Persist LLM outputs to HIEROMEM."""
        if not self.persist_outputs:
            return

        for generation in response.generations:
            for gen in generation:
                self._insert_episode(
                    content=gen.text,
                    meta={
                        "type": "llm_output",
                        "run_id": str(run_id),
                    },
                )

    def on_chain_start(
        self,
        serialized: Dict[str, Any],
        inputs: Dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Persist chain inputs."""
        if not self.persist_inputs:
            return

        # Serialize inputs to string
        input_str = str(inputs)
        self._insert_episode(
            content=input_str,
            meta={
                "type": "chain_input",
                "run_id": str(run_id),
                "chain_name": serialized.get("name", "unknown"),
            },
        )

    def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Persist chain outputs."""
        if not self.persist_outputs:
            return

        output_str = str(outputs)
        self._insert_episode(
            content=output_str,
            meta={
                "type": "chain_output",
                "run_id": str(run_id),
            },
        )

    def _insert_episode(self, content: str, meta: Dict[str, Any]) -> None:
        """Insert episode to HIEROMEM."""
        if self._client is None:
            return

        try:
            self._client.post(
                "/memory/insert",
                json={
                    "text": content,
                    "namespace": self.namespace,
                    "meta": meta,
                },
            )
        except httpx.HTTPError:
            pass  # Silent fail

    def __del__(self) -> None:
        if hasattr(self, "_client") and self._client:
            self._client.close()
