"""LangChain adapter — a LangChain-compatible memory class backed by LociDB."""

from __future__ import annotations

import re
from typing import Any, Dict, List, Optional

from ..core.db import LociDB


def _extract_entities(text: str) -> List[str]:
    entities: set[str] = set()
    for m in re.finditer(r'"([^"]+)"', text):
        entities.add(m.group(1))
    for m in re.finditer(r"\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\b", text):
        word = m.group(1)
        if len(word) > 2 and word not in {"The", "This", "That", "What", "How", "Why", "When", "Where"}:
            entities.add(word)
    return list(entities)


class LociLangChainMemory:
    """A LangChain-compatible memory backend powered by LociDB.

    Works as a drop-in ``memory`` for LangChain chains and agents.

    Usage::

        from locidb.adapters import LociLangChainMemory
        from langchain.chains import ConversationChain
        from langchain_openai import ChatOpenAI

        memory = LociLangChainMemory(db_path="./memory.loci")
        chain = ConversationChain(llm=ChatOpenAI(), memory=memory)
        chain.run("Tell me about Dropzi")
    """

    memory_key: str = "history"
    return_messages: bool = False

    def __init__(
        self,
        db_path: str = "./memory.loci",
        scope: str = "global",
        recall_top_k: int = 8,
        min_importance: float = 0.3,
        time_window_days: Optional[int] = 30,
    ) -> None:
        self.db = LociDB(path=db_path)
        self.scope = scope
        self.recall_top_k = recall_top_k
        self.min_importance = min_importance
        self.time_window_days = time_window_days
        self._buffer: List[Dict[str, str]] = []

    @property
    def memory_variables(self) -> List[str]:
        return [self.memory_key]

    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Load relevant memories based on the current input."""
        query = ""
        for key in ("input", "question", "query", "human_input"):
            if key in inputs:
                query = str(inputs[key])
                break

        if not query:
            return {self.memory_key: "" if not self.return_messages else []}

        entities = _extract_entities(query)
        results = self.db.recall(
            query=query,
            entities=entities or None,
            scope=self.scope,
            time_window_days=self.time_window_days,
            min_importance=self.min_importance,
            top_k=self.recall_top_k,
        )

        # Build context with persona + corrections + memories
        persona_prompt = self.db.get_persona_prompt()
        corrections = self.db.get_corrections(scope=self.scope, limit=5)

        if self.return_messages:
            messages = []
            if persona_prompt:
                messages.append({"role": "system", "content": persona_prompt})
            for c in corrections:
                messages.append({"role": "system", "content": f"[Correction] {c.content}"})
            for r in results:
                messages.append({
                    "role": "system",
                    "content": f"[Memory] {r['node'].content}",
                })
            return {self.memory_key: messages}
        else:
            lines = []
            if persona_prompt:
                lines.append(persona_prompt)
            if corrections:
                lines.append("Corrections:")
                for c in corrections:
                    lines.append(f"- {c.content}")
            lines.extend(f"- {r['node'].content}" for r in results)
            return {self.memory_key: "\n".join(lines)}

    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
        """Save the current exchange as memories."""
        # Save user input
        user_text = ""
        for key in ("input", "question", "query", "human_input"):
            if key in inputs:
                user_text = str(inputs[key])
                break
        if user_text:
            entities = _extract_entities(user_text)
            self.db.memorize(content=user_text, entities=entities, scope=self.scope)

        # Save assistant output
        for key in ("output", "response", "answer", "text"):
            if key in outputs:
                assistant_text = str(outputs[key])
                entities = _extract_entities(assistant_text)
                self.db.memorize(content=assistant_text, entities=entities, scope=self.scope)
                break

    def clear(self) -> None:
        """Clear the in-memory buffer (does not delete persisted memories)."""
        self._buffer.clear()

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