backend/routes/memory_routes.py

from fastapi import HTTPException, APIRouter
from pydantic import BaseModel
from typing import Optional

router = APIRouter()

class MemoryItem(BaseModel):
    content: str
    metadata: Optional[dict] = None
    persona: Optional[str] = "jarvis"

class SearchQuery(BaseModel):
    query: str
    persona: Optional[str] = "jarvis"

def get_memory_client():
    from backend.memory.episodic_memory import EpisodicMemory
    import os, sys
    db_dir = os.path.join(os.environ.get('APPDATA', os.path.expanduser('~')), 'JARVIS_OS') if getattr(sys, 'frozen', False) else os.path.join(os.getcwd(), 'chroma_db')
    os.makedirs(db_dir, exist_ok=True)
    return EpisodicMemory(db_dir)

@router.get("/list")
async def list_memories(persona: Optional[str] = "jarvis"):
    mem = get_memory_client()
    if mem and mem.client:
        results = await mem.query("memory", top_k=50, persona=persona)
        return results
    return []

@router.post("/search")
async def search_memories(q: SearchQuery):
    mem = get_memory_client()
    if mem and mem.client:
        results = await mem.query(q.query, top_k=20, persona=q.persona)
        return results
    return []

@router.post("/store")
async def store_memory(item: MemoryItem):
    if len(item.content.encode("utf-8")) > 100 * 1024:
        raise HTTPException(status_code=400, detail="Payload exceeds 100KB memory limit.")
    mem = get_memory_client()
    if mem and mem.client:
        await mem.add(item.content, item.metadata or {"source": "exe_rest"}, item.persona)
    return {"status": "stored"}

@router.delete("/{id}")
async def delete_memory(id: str, persona: Optional[str] = "jarvis"):
    mem = get_memory_client()
    if mem and mem.client:
        await mem.delete(id, persona)
    return {"status": "deleted"}

@router.get("/graph")
async def get_graph():
    from backend.memory.semantic_memory import SemanticMemory
    try:
        mem = SemanticMemory()
        nodes = []
        edges = []
        
        async for node in mem.nodes.find({}):
            nodes.append({"id": node["_id"], "label": node["label"], "properties": str(node.get("properties", {}))})
            
        async for edge in mem.edges.find({}):
            edges.append({"source": edge["from_id"], "target": edge["to_id"], "label": edge["relation"], "weight": edge.get("weight", 1.0)})
            
        return {"nodes": nodes, "edges": edges}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"MongoDB Atlas Error: {str(e)}")

@router.get("/history")
async def get_history(limit: int = 50, offset: int = 0, persona: Optional[str] = None):
    import os, sqlite3
    from backend.services.usb_monitor import get_db_path
    db_path = get_db_path()
        
    history = []
    if os.path.exists(db_path):
        try:
            with sqlite3.connect(db_path) as conn:
                conn.execute('PRAGMA journal_mode=WAL')
                # First ensure persona_id exists so queries don't crash
                try:
                    conn.execute("ALTER TABLE conversations ADD COLUMN persona_id TEXT DEFAULT 'jarvis'")
                except sqlite3.OperationalError:
                    pass
                
                cursor = conn.cursor()
                if persona:
                    cursor.execute("SELECT id, role, content, timestamp FROM conversations WHERE persona_id = ? ORDER BY timestamp DESC LIMIT ? OFFSET ?", (persona.lower(), limit, offset))
                else:
                    cursor.execute("SELECT id, role, content, timestamp FROM conversations ORDER BY timestamp DESC LIMIT ? OFFSET ?", (limit, offset))
                for row in cursor.fetchall():
                    history.append({
                        "id": row[0],
                        "sender": "agent" if row[1] == "assistant" else "user",
                        "text": row[2],
                        "timestamp": str(row[3])
                    })
            history.reverse()  # UI needs chronological order (oldest first)
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))
    return history