import json
import os
from collections import deque
from pathlib import Path
from typing import Any, Dict, List, Optional
from .config import WORKING_MEMORY_SIZE, MEMORY_DIR

class MemorySystem:
    def __init__(self):
        self.working = deque(maxlen=WORKING_MEMORY_SIZE)
        self._long_term: List[Dict[str, Any]] = []
        self._load_long_term()

    def _long_term_path(self) -> Path:
        return MEMORY_DIR / "long_term.jsonl"

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

    def _save_long_term(self):
        path = self._long_term_path()
        with open(path, "w") as f:
            for entry in self._long_term[-1000:]:
                f.write(json.dumps(entry) + "\n")

    def add_working(self, entry: Dict[str, Any]):
        self.working.append(entry)

    def get_working(self) -> List[Dict[str, Any]]:
        return list(self.working)

    def add_long_term(self, entry: Dict[str, Any]):
        self._long_term.append(entry)
        self._save_long_term()

    def query_long_term(self, key: str, top_k: int = 5) -> List[Dict[str, Any]]:
        matches = []
        for entry in reversed(self._long_term):
            if key.lower() in str(entry).lower():
                matches.append(entry)
                if len(matches) >= top_k:
                    break
        return matches

    def get_context(self, query: str) -> str:
        working_summary = "\n".join(
            f"{e.get('role', 'system')}: {e.get('content', '')[:200]}"
            for e in self.working[-5:]
        )
        lt_matches = self.query_long_term(query, top_k=3)
        lt_summary = "\n".join(
            f"[memory] {m.get('summary', str(m)[:200])}" for m in lt_matches
        )
        parts = []
        if working_summary:
            parts.append(f"[working memory]\n{working_summary}")
        if lt_summary:
            parts.append(f"[long-term memory]\n{lt_summary}")
        return "\n\n".join(parts)

    def summarize_to_long_term(self, role: str, content: str):
        summary = content[:300] if len(content) > 300 else content
        self.add_long_term({"role": role, "summary": summary, "content": content[:1000]})