import os
import json
import time
from typing import List, Dict, Any, Optional

class FlatFileMemoryManager:
    """
    SOTA Lightweight Flat-File Memory Manager.
    Stores conversation history in per-user JSONL files.
    Optimized for memory and storage efficiency.
    """
    def __init__(self, memory_dir: str = "memories", max_history: int = 100):
        self.memory_dir = memory_dir
        self.max_history = max_history
        os.makedirs(self.memory_dir, exist_ok=True)

    def get_user_file_path(self, user_id: str) -> str:
        return os.path.join(self.memory_dir, f"{user_id}.jsonl")

    def add_interaction(self, user_id: str, prompt: str, response: str, modality: str = "text"):
        """Appends an interaction to the user's JSONL memory file."""
        file_path = self.get_user_file_path(user_id)
        timestamp = time.time()
        
        interaction = {
            "timestamp": timestamp,
            "turns": [
                {"role": "user", "content": prompt, "modality": modality},
                {"role": "assistant", "content": response, "modality": modality}
            ]
        }
        
        with open(file_path, "a") as f:
            f.write(json.dumps(interaction) + "\n")

    def get_context(self, user_id: str, limit: Optional[int] = None) -> List[Dict[str, str]]:
        """Retrieves recent interaction history using an efficient tail-read."""
        limit = limit or self.max_history
        file_path = self.get_user_file_path(user_id)
        
        if not os.path.exists(file_path):
            return []

        # Efficiently read last N lines
        lines: List[str] = []
        try:
            with open(file_path, "r") as f:
                all_lines = f.readlines()
                final_limit: int = limit if isinstance(limit, int) else self.max_history
                start_idx = max(0, len(all_lines) - final_limit)
                lines = all_lines[start_idx:]
        except Exception as e:
            print(f"[Memory] Error reading memory for {user_id}: {e}")
            return []

        history = []
        for line in lines:
            try:
                data = json.loads(line)
                for turn in data["turns"]:
                    history.append({
                        "role": turn["role"],
                        "content": turn["content"],
                        "modality_hint": turn.get("modality", "text")
                    })
            except Exception:
                continue
        return history

    def get_training_samples(self, user_id: str, limit: int = 50) -> List[Dict[str, str]]:
        """Retrieves recent interactions formatted for on-the-fly training."""
        history = self.get_context(user_id, limit=limit)
        samples = []
        for i in range(0, len(history) - 1, 2):
            if i + 1 < len(history):
                samples.append({
                    "instruction": history[i]["content"],
                    "response": history[i+1]["content"]
                })
        return samples

    def clear_memory(self, user_id: str):
        """Wipe a user's memory file."""
        file_path = self.get_user_file_path(user_id)
        if os.path.exists(file_path):
            os.remove(file_path)
