# backend/memory/consolidator.py
import os
import json
import logging

from .working_memory import WorkingMemory
from .episodic_memory import EpisodicMemory
from .semantic_memory import SemanticMemory
from .procedural_memory import ProceduralMemory

try:
    from google import genai
    from google.genai import types
    from config import GEMINI_API_KEY, GEMINI_MODEL
except ImportError:
    pass

class MemoryConsolidator:
    def __init__(self):
        # Resolve persistence directory from environment variable provided by Tauri
        base_dir = os.environ.get("JARVIS_APP_DATA_DIR")
        if not base_dir:
            # Fallback to local if not set
            base_dir = os.path.join(os.path.expanduser("~"), ".jarvis_omega")
            
        mem_dir = os.path.join(base_dir, "memory")
        os.makedirs(mem_dir, exist_ok=True)
        
        self.episodic = EpisodicMemory(os.path.join(mem_dir, "episodic"))
        self.semantic = SemanticMemory(os.path.join(mem_dir, "semantic.db"))
        self.procedural = ProceduralMemory(os.path.join(mem_dir, "procedural.db"))
        
        self._client = None
        try:
            self._client = genai.Client(api_key=GEMINI_API_KEY)
        except Exception as e:
            import logging; logging.getLogger(__name__).error(f"Swallowed exception: {e}")

    async def _ask_llm_json(self, prompt: str, system: str) -> dict:
        if not self._client:
            return {}
        try:
            resp = await self._client.aio.models.generate_content(
                model=GEMINI_MODEL,
                contents=prompt,
                config=types.GenerateContentConfig(
                    system_instruction=system,
                    temperature=0.1,
                    response_mime_type="application/json"
                ),
            )
            raw = resp.text.strip()
            if raw.startswith("```"):
                import re
                raw = re.sub(r"^```(?:json)?\s*|\s*```$", "", raw, flags=re.MULTILINE)
            return json.loads(raw)
        except Exception as e:
            logging.error(f"Consolidator LLM Error: {e}")
            return {}

    async def consolidate(self, working_memory: WorkingMemory):
        """Runs at session end to consolidate working memory."""
        if not working_memory.messages:
            return
            
        # 1. Summarize into a single episode
        convo_text = ""
        for msg in working_memory.messages:
            convo_text += f"{msg['role'].upper()}: {msg['content']}\n"
            
        summary_prompt = f"Summarize this conversation into a single descriptive paragraph:\n\n{convo_text}"
        summary_sys = 'You are a memory consolidation engine. Return a JSON object: {"summary": "..."}'
        
        summary_result = await self._ask_llm_json(summary_prompt, summary_sys)
        episode_summary = summary_result.get("summary", "Conversation recorded.")
        
        # Call EpisodicMemory.add()
        await self.episodic.add(episode_summary, {"type": "session_summary", "msg_count": len(working_memory.messages)})
        
        # 2. Extract semantic triples
        triple_prompt = f"Extract key factual knowledge from this conversation as subject-predicate-object triples.\nConversation:\n{convo_text}"
        triple_sys = """Return a JSON array of triples:
{
  "triples": [
    {"subject": "Tony Stark", "predicate": "likes", "object": "cheeseburgers", "confidence": 0.9}
  ]
}"""
        triple_result = await self._ask_llm_json(triple_prompt, triple_sys)
        triples = triple_result.get("triples", [])
        
        for t in triples:
            if "subject" in t and "predicate" in t and "object" in t:
                await self.semantic.add_fact(
                    subject=t["subject"],
                    predicate=t["predicate"],
                    obj=t["object"],
                    confidence=t.get("confidence", 1.0)
                )
                
        # 3. Extract procedural patterns (if any new tasks were successfully completed)
        proc_prompt = f"Did the AI successfully complete any multi-step task in this conversation? If so, extract the pattern.\nConversation:\n{convo_text}"
        proc_sys = """Return JSON format:
{
  "tasks": [
    {
      "task_name": "clear_temp_files",
      "steps": [{"action": "run_shell_command", "args": {"command": "rm -rf /tmp/*"}}],
      "triggers": ["clear temp", "empty trash"]
    }
  ]
}
If none, return {"tasks": []}"""
        proc_result = await self._ask_llm_json(proc_prompt, proc_sys)
        tasks = proc_result.get("tasks", [])
        
        for t in tasks:
            if "task_name" in t and "steps" in t:
                # Update ProceduralMemory (async — must be awaited or the Mongo
                # write is silently dropped, same class of bug as add_fact above).
                await self.procedural.record_success(t["task_name"], t["steps"])
                
        # Finally, clear working memory
        working_memory.clear()
        logging.info("Memory consolidation complete.")