import os
import time
import hashlib
from schemas.agent import MemoryEntry

DB_PATH = os.getenv("AGENT_MEMORY_PATH", "/tmp/agent_memory")


class AgentMemory:
    def __init__(self):
        self._store = {}
        self._init_path()

    def _init_path(self):
        os.makedirs(DB_PATH, exist_ok=True)

    async def store(self, entry: MemoryEntry) -> str:
        entry_id = hashlib.md5(f"{entry.session_id}:{entry.action}:{time.time()}".encode()).hexdigest()
        key = f"{entry.session_id}:{entry_id}"
        self._store[key] = {
            "session_id": entry.session_id,
            "task_id": entry.task_id,
            "action": entry.action,
            "prompt": entry.prompt,
            "response": entry.response,
            "tool_used": entry.tool_used,
            "result": entry.result,
            "success": entry.success,
            "tokens_used": entry.tokens_used,
            "latency_ms": entry.latency_ms,
            "created_at": time.time()
        }
        return entry_id

    async def recall(self, session_id: str, query: str = "", limit: int = 20) -> list[MemoryEntry]:
        results = []
        for k, v in sorted(self._store.items(), key=lambda x: x[1].get("created_at", 0), reverse=True):
            if v["session_id"] == session_id:
                results.append(MemoryEntry(
                    session_id=v["session_id"],
                    task_id=v.get("task_id"),
                    action=v.get("action", ""),
                    prompt=v.get("prompt"),
                    response=v.get("response"),
                    tool_used=v.get("tool_used"),
                    result=v.get("result"),
                    success=v.get("success", True),
                    tokens_used=v.get("tokens_used", 0),
                    latency_ms=v.get("latency_ms", 0)
                ))
                if len(results) >= limit:
                    break
        if query and results:
            scored = [(self._simple_score(entry, query), entry) for entry in results]
            scored.sort(key=lambda x: x[0], reverse=True)
            results = [e for _, e in scored[:limit]]
        return results

    def _simple_score(self, entry: MemoryEntry, query: str) -> float:
        q = query.lower()
        score = 0.0
        if entry.action and q in entry.action.lower():
            score += 0.5
        if entry.result and q in entry.result.lower():
            score += 0.4
        if entry.prompt and q in entry.prompt.lower():
            score += 0.3
        if entry.response and q in entry.response.lower():
            score += 0.2
        if entry.tool_used and q in entry.tool_used.lower():
            score += 0.1
        return score

    async def get_session_summary(self, session_id: str) -> str:
        entries = await self.recall(session_id, "", limit=50)
        if not entries:
            return "No prior session data."
        success_count = sum(1 for e in entries if e.success)
        total = len(entries)
        tools_used = list(set(e.tool_used for e in entries if e.tool_used))
        last_action = entries[0].action if entries else "none"
        return (f"Session: {success_count}/{total} successful actions. "
                f"Tools used: {', '.join(tools_used) if tools_used else 'none'}. "
                f"Last action: {last_action}")

    async def clear_session(self, session_id: str):
        keys_to_del = [k for k in self._store if k.startswith(f"{session_id}:")]
        for k in keys_to_del:
            self._store.pop(k, None)