"""
Hierarchical Memory — Episodic, Semantic, and Procedural tiers.
Episodic: raw event log (time-indexed)
Semantic: knowledge graph of Atomic Truths (distilled from episodic)
Procedural: learned skills and protocols
"""

import json
import time
from collections import defaultdict
from pathlib import Path
from typing import Any

from config.config import MEMORY_DIR


class EpisodicMemory:
    def __init__(self, agent_id: str = "vitalis"):
        self.path = MEMORY_DIR / f"{agent_id}_episodic.jsonl"
        self._events: list[dict] = []
        self._load()

    def _load(self):
        if self.path.exists():
            with open(self.path) as f:
                for line in f:
                    line = line.strip()
                    if line:
                        self._events.append(json.loads(line))

    def record(self, event_type: str, content: dict[str, Any],
               source: str = "perception"):
        entry = {
            "id": f"ep-{int(time.time() * 1000)}-{len(self._events)}",
            "timestamp": time.time(),
            "type": event_type,
            "source": source,
            "content": content,
        }
        self._events.append(entry)
        with open(self.path, "a") as f:
            f.write(json.dumps(entry) + "\n")
        return entry

    def query(self, query_type: str | None = None,
              limit: int = 50, since: float = 0) -> list[dict]:
        results = []
        for e in reversed(self._events):
            if e["timestamp"] < since:
                continue
            if query_type and e["type"] != query_type:
                continue
            results.append(e)
            if len(results) >= limit:
                break
        return results

    def recent(self, seconds: float = 300) -> list[dict]:
        cutoff = time.time() - seconds
        return self.query(since=cutoff)

    def count(self) -> int:
        return len(self._events)


class AtomicTruth:
    def __init__(self, truth_id: str, statement: str, confidence: float,
                 source_events: list[str], category: str = "general"):
        self.truth_id = truth_id
        self.statement = statement
        self.confidence = confidence
        self.source_events = source_events
        self.category = category
        self.created_at = time.time()
        self.access_count = 0

    def to_dict(self) -> dict:
        return {
            "truth_id": self.truth_id,
            "statement": self.statement,
            "confidence": self.confidence,
            "source_events": self.source_events,
            "category": self.category,
            "created_at": self.created_at,
            "access_count": self.access_count,
        }

    @classmethod
    def from_dict(cls, d: dict) -> "AtomicTruth":
        t = cls(d["truth_id"], d["statement"], d["confidence"],
                d["source_events"], d.get("category", "general"))
        t.created_at = d.get("created_at", time.time())
        t.access_count = d.get("access_count", 0)
        return t


class SemanticMemory:
    def __init__(self, agent_id: str = "vitalis"):
        self.path = MEMORY_DIR / f"{agent_id}_semantic.json"
        self._truths: dict[str, AtomicTruth] = {}
        self._load()

    def _load(self):
        if self.path.exists():
            data = json.loads(self.path.read_text())
            for d in data:
                truth = AtomicTruth.from_dict(d)
                self._truths[truth.truth_id] = truth

    def _save(self):
        with open(self.path, "w") as f:
            json.dump([t.to_dict() for t in self._truths.values()], f, indent=2)

    def add(self, statement: str, confidence: float,
            source_events: list[str], category: str = "general") -> AtomicTruth:
        tid = f"at-{int(time.time())}-{hash(statement) % 10000}"
        truth = AtomicTruth(tid, statement, confidence, source_events, category)
        self._truths[tid] = truth
        self._save()
        return truth

    def query(self, text: str, top_k: int = 5) -> list[AtomicTruth]:
        text_lower = text.lower()
        scored = []
        for t in self._truths.values():
            score = 0.0
            if text_lower in t.statement.lower():
                score += t.confidence * 0.8
            words = set(text_lower.split())
            truth_words = set(t.statement.lower().split())
            overlap = len(words & truth_words) / max(len(words), 1)
            score += overlap * 0.2
            scored.append((score, t))
        scored.sort(key=lambda x: -x[0])
        results = [t for _, t in scored[:top_k]]
        for t in results:
            t.access_count += 1
        self._save()
        return results

    def get_by_category(self, category: str) -> list[AtomicTruth]:
        return [t for t in self._truths.values() if t.category == category]

    def count(self) -> int:
        return len(self._truths)


class ProceduralMemory:
    def __init__(self, agent_id: str = "vitalis"):
        self.path = MEMORY_DIR / f"{agent_id}_procedural.json"
        self._skills: dict[str, dict] = {}
        self._load()

    def _load(self):
        if self.path.exists():
            self._skills = json.loads(self.path.read_text())

    def _save(self):
        with open(self.path, "w") as f:
            json.dump(self._skills, f, indent=2)

    def record_skill(self, name: str, description: str,
                     protocol: list[str], success_rate: float = 1.0):
        self._skills[name] = {
            "description": description,
            "protocol": protocol,
            "success_rate": success_rate,
            "use_count": 0,
            "created_at": time.time(),
        }
        self._save()

    def retrieve(self, task: str) -> list[tuple[str, dict]]:
        task_lower = task.lower()
        results = []
        for name, skill in self._skills.items():
            if task_lower in name.lower() or task_lower in skill.get("description", "").lower():
                results.append((name, skill))
        results.sort(key=lambda x: -x[1].get("success_rate", 0))
        return results

    def record_success(self, name: str):
        skill = self._skills.get(name)
        if skill:
            skill["use_count"] = skill.get("use_count", 0) + 1
            self._save()

    def record_failure(self, name: str):
        skill = self._skills.get(name)
        if skill:
            skill["success_rate"] = max(0.0, skill.get("success_rate", 1.0) - 0.1)
            self._save()

    def count(self) -> int:
        return len(self._skills)