"""
Synrix Agent Runtime — OpenAI Agents SDK Integration
Persistent memory for OpenAI Agents SDK.

Usage:
    from synrix_runtime.integrations.openai_agents import SynrixOpenAIMemory
    memory = SynrixOpenAIMemory()
    memory.store_thread_state("thread_123", {"messages": [...], "context": {...}})
"""

import time
import json
from typing import Dict, List, Optional, Any
from synrix.agent_backend import get_synrix_backend


class SynrixOpenAIMemory:
    """Persistent memory for OpenAI Agents SDK."""

    def __init__(self, backend=None):
        if backend is not None:
            self.backend = backend
        else:
            from synrix_runtime.config import SynrixConfig
            config = SynrixConfig.from_env()
            self.backend = get_synrix_backend(**config.get_backend_kwargs())

    def store_thread_state(self, thread_id: str, state: dict):
        """Store the state of a thread."""
        payload = state.copy()
        payload["_stored_at"] = time.time()

        start = time.perf_counter_ns()
        self.backend.write(
            f"openai:threads:{thread_id}:state",
            payload,
            metadata={"type": "openai_thread_state", "thread_id": thread_id}
        )
        latency_us = (time.perf_counter_ns() - start) / 1000
        return {"thread_id": thread_id, "latency_us": latency_us}

    def restore_thread(self, thread_id: str) -> Optional[dict]:
        """Restore a thread's state from Synrix."""
        start = time.perf_counter_ns()
        result = self.backend.read(f"openai:threads:{thread_id}:state")
        latency_us = (time.perf_counter_ns() - start) / 1000

        if result:
            data = result.get("data", {})
            val = data.get("value", data)
            return {"state": val, "latency_us": latency_us, "found": True}
        return {"state": None, "latency_us": latency_us, "found": False}

    def store_run_result(self, run_id: str, result: dict):
        """Store the result of an agent run."""
        payload = result.copy() if isinstance(result, dict) else {"value": result}
        payload["_stored_at"] = time.time()

        start = time.perf_counter_ns()
        self.backend.write(
            f"openai:runs:{run_id}",
            payload,
            metadata={"type": "openai_run_result", "run_id": run_id}
        )
        latency_us = (time.perf_counter_ns() - start) / 1000
        return {"run_id": run_id, "latency_us": latency_us}

    def get_agent_history(self, agent_name: str) -> list:
        """Get all runs for a specific agent."""
        results = self.backend.query_prefix(f"openai:runs:", limit=200)
        history = []
        for r in results:
            data = r.get("data", {})
            val = data.get("value", data)
            if isinstance(val, dict) and val.get("agent_name") == agent_name:
                history.append(val)
        history.sort(key=lambda x: x.get("_stored_at", 0), reverse=True)
        return history

    def get_all_threads(self) -> list:
        """List all stored threads."""
        results = self.backend.query_prefix("openai:threads:", limit=200)
        threads = []
        seen = set()
        for r in results:
            key = r.get("key", "")
            parts = key.split(":")
            if len(parts) >= 3:
                thread_id = parts[2]
                if thread_id not in seen:
                    seen.add(thread_id)
                    data = r.get("data", {})
                    val = data.get("value", data)
                    threads.append({"thread_id": thread_id, "state": val})
        return threads

    def get_all_runs(self) -> list:
        """List all stored runs."""
        results = self.backend.query_prefix("openai:runs:", limit=200)
        runs = []
        for r in results:
            data = r.get("data", {})
            val = data.get("value", data)
            runs.append(val)
        runs.sort(key=lambda x: x.get("_stored_at", 0) if isinstance(x, dict) else 0, reverse=True)
        return runs
