import os
import datetime
import chromadb

DEFAULT_DB_PATH = os.path.expanduser(
    "/home/oem/CyberBrain26/workspace/chroma_db"
)


class CyberBrainMemory:
    """
    Tenhle wrapper pracuje přímo s existující ChromaDB, kterou už používáš.
    Nezasahuje do tvého původního MemoryManageru, jen ji doplňuje o:
    - uloženou identitu CyberBrain26
    - helper na získání kontextu pro system prompt
    """

    def __init__(self, persist_directory: str = DEFAULT_DB_PATH):
        self.persist_directory = persist_directory
        self.client = chromadb.PersistentClient(path=persist_directory)

        # Kolekce, které už máš ve workspace/chroma_db
        self.preferences = self._get_or_create_collection("preferences")
        self.rules = self._get_or_create_collection("rules")
        self.conversations = self._get_or_create_collection("conversations")
        self.projects = self._get_or_create_collection("projects")
        self.errors = self._get_or_create_collection("errors")

        # Při startu zajistíme identitu
        self._ensure_identity_rule()

    def _get_or_create_collection(self, name: str):
        try:
            return self.client.get_collection(name)
        except Exception:
            return self.client.create_collection(name=name)

    def _now_iso(self) -> str:
        return datetime.datetime.now().isoformat()

    def _ensure_identity_rule(self):
        """
        Uloží (nebo aktualizuje) pravidlo s identitou CyberBrain26.
        Díky tomu je to vždy v paměti bez ohledu na konkrétní system prompt.
        """

        doc = (
            "Jsi lokální AI systém CyberBrain26 běžící na Linux Mint. "
            "Řídíš se principy: bezpečnost, transparentnost, reprodukovatelnost "
            "a dlouhodobá paměť. Mluvíš česky, ale rozumíš i anglicky."
        )

        # Jednoduché ID, ať se ti nevytváří stovky duplikátů
        identity_id = "rule-identity-cyberbrain26"

        self.rules.upsert(
            ids=[identity_id],
            documents=[doc],
            metadatas=[{"type": "identity", "timestamp": self._now_iso()}],
        )

    def save_preference(self, text: str, meta: dict | None = None):
        base_meta = {"type": "preference", "timestamp": self._now_iso()}
        if meta:
            base_meta.update(meta)
        self.preferences.add(
            ids=[f"pref-{self._now_iso()}"],
            documents=[text],
            metadatas=[base_meta],
        )

    def save_rule(self, text: str, meta: dict | None = None):
        base_meta = {"type": "rule", "timestamp": self._now_iso()}
        if meta:
            base_meta.update(meta)
        self.rules.add(
            ids=[f"rule-{self._now_iso()}"],
            documents=[text],
            metadatas=[base_meta],
        )

    def save_conversation_turn(
        self,
        user_text: str,
        assistant_text: str,
        meta: dict | None = None,
    ):
        """
        Jednoduché uložení jednoho „turnu“ konverzace.
        Můžeš volat vždy, když safe_oi_runner skončí odpověď.
        """

        base_meta = {
            "type": "conversation",
            "timestamp": self._now_iso(),
        }
        if meta:
            base_meta.update(meta)

        doc = f"User: {user_text}\nAssistant: {assistant_text}"

        self.conversations.add(
            ids=[f"conv-{self._now_iso()}"],
            documents=[doc],
            metadatas=[base_meta],
        )

    def retrieve_context(self, user_query: str, n_results: int = 5) -> str:
        """
        Vrátí textový kontext složený z:
        - identity (rules: type=identity)
        - preferencí (preferences)
        - pravidel (rules)
        - relevantních konverzací podle user_query

        Tenhle text pak přilepíš k system promptu.
        """

        parts: list[str] = []

        # 1) Identita
        try:
            identity = self.rules.get(
                ids=["rule-identity-cyberbrain26"],
            )
            if identity and identity.get("documents"):
                parts.append("IDENTITA:\n" + identity["documents"][0])
        except Exception:
            # když ještě nic není, prostě pokračujeme
            pass

        # 2) Preference – vezmeme pár posledních
        try:
            prefs = self.preferences.get(limit=3)
            docs = prefs.get("documents") or []
            if docs:
                text = "\n".join(docs)
                parts.append("PREFERENCE:\n" + text)
        except Exception:
            pass

        # 3) Pravidla – mimo identity vezmeme pár posledních
        try:
            rules = self.rules.get(limit=5)
            docs = rules.get("documents") or []
            other_rules = [
                d for d in docs if "CyberBrain26" not in d
            ]
            if other_rules:
                text = "\n".join(other_rules)
                parts.append("PRAVIDLA:\n" + text)
        except Exception:
            pass

        # 4) Relevantní konverzace podle dotazu
        try:
            conv_query = self.conversations.query(
                query_texts=[user_query],
                n_results=n_results,
            )
            conv_docs = conv_query.get("documents") or [[]]
            flat = []
            for lst in conv_docs:
                flat.extend(lst)
            if flat:
                text = "\n---\n".join(flat)
                parts.append("RELEVANTNÍ KONVERZACE:\n" + text)
        except Exception:
            pass

        if not parts:
            return ""

        return "\n\n=== KONTEKST Z PAMĚTI ===\n" + "\n\n".join(parts)


def build_system_prompt_with_memory(
    base_system_prompt: str,
    user_query: str,
    memory: CyberBrainMemory | None = None,
) -> str:
    """
    Helper, který můžeš přímo použít v safe_oi_runner.py:

        mem = CyberBrainMemory()
        system_prompt = build_system_prompt_with_memory(
            BASE_SYSTEM_PROMPT,
            user_query,
            mem,
        )

    Vrátí base_system_prompt doplněný o kontext z paměti.
    """

    if memory is None:
        memory = CyberBrainMemory()

    context = memory.retrieve_context(user_query=user_query, n_results=5)

    if not context:
        return base_system_prompt

    return base_system_prompt + "\n\n" + context