"""
Synrix Agent Runtime — CrewAI Integration
Shared persistent memory for CrewAI crews.

Usage:
    from synrix_runtime.integrations.crewai_memory import SynrixCrewMemory
    crew_memory = SynrixCrewMemory(crew_id="research_crew")
"""

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


class SynrixCrewMemory:
    """
    Shared persistent memory for CrewAI crews.
    All agents in the crew share knowledge through Synrix.

    Usage:
        crew_memory = SynrixCrewMemory(crew_id="research_crew")
        crew_memory.store_finding("researcher", "market_size", {"value": "$4.2B"})
        finding = crew_memory.get_finding("market_size")
    """

    def __init__(self, crew_id: str, backend=None):
        self.crew_id = crew_id
        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())

        start = time.perf_counter_ns()
        self.backend.write(
            f"crewai:{crew_id}:meta",
            {"crew_id": crew_id, "created_at": time.time()},
            metadata={"type": "crew_meta"}
        )
        latency_us = (time.perf_counter_ns() - start) / 1000
        print(f"[CrewAI] Crew '{crew_id}' connected to Synrix in {latency_us:.1f}us")

    def store_finding(self, agent_role: str, key: str, finding: Any):
        """Store a finding from a crew agent."""
        payload = finding if isinstance(finding, dict) else {"value": finding}
        payload["_agent_role"] = agent_role
        payload["_stored_at"] = time.time()

        start = time.perf_counter_ns()
        self.backend.write(
            f"crewai:{self.crew_id}:findings:{key}",
            payload,
            metadata={"type": "crew_finding", "agent_role": agent_role}
        )
        latency_us = (time.perf_counter_ns() - start) / 1000
        return {"key": key, "latency_us": latency_us}

    def get_finding(self, key: str) -> Optional[dict]:
        """Get a specific finding."""
        result = self.backend.read(f"crewai:{self.crew_id}:findings:{key}")
        if result:
            data = result.get("data", {})
            return data.get("value", data)
        return None

    def get_all_findings(self) -> list:
        """Get all findings from the crew."""
        results = self.backend.query_prefix(f"crewai:{self.crew_id}:findings:", limit=200)
        findings = []
        for r in results:
            key = r.get("key", "").replace(f"crewai:{self.crew_id}:findings:", "")
            data = r.get("data", {})
            val = data.get("value", data)
            findings.append({"key": key, "data": val})
        return findings

    def store_task_result(self, task_name: str, result: Any, agent_role: str):
        """Store a task result."""
        payload = result if isinstance(result, dict) else {"value": result}
        payload["_agent_role"] = agent_role
        payload["_completed_at"] = time.time()

        self.backend.write(
            f"crewai:{self.crew_id}:tasks:{task_name}",
            payload,
            metadata={"type": "crew_task_result", "agent_role": agent_role}
        )

    def get_crew_knowledge_base(self) -> dict:
        """Get the entire crew knowledge base."""
        findings = self.get_all_findings()
        tasks = self.backend.query_prefix(f"crewai:{self.crew_id}:tasks:", limit=200)

        task_results = []
        for t in tasks:
            key = t.get("key", "").replace(f"crewai:{self.crew_id}:tasks:", "")
            data = t.get("data", {})
            val = data.get("value", data)
            task_results.append({"task": key, "result": val})

        return {
            "crew_id": self.crew_id,
            "findings": findings,
            "task_results": task_results,
            "total_items": len(findings) + len(task_results),
        }

    def crew_snapshot(self, label: str = None):
        """Snapshot entire crew state."""
        if label is None:
            label = f"crew_snap_{int(time.time()*1000000)}"

        kb = self.get_crew_knowledge_base()
        start = time.perf_counter_ns()
        self.backend.write(
            f"crewai:{self.crew_id}:snapshots:{label}",
            {"label": label, "knowledge_base": kb, "created_at": time.time()},
            metadata={"type": "crew_snapshot"}
        )
        latency_us = (time.perf_counter_ns() - start) / 1000
        return {"label": label, "items": kb["total_items"], "latency_us": latency_us}

    def crew_restore(self, label: str) -> dict:
        """Restore crew state from a snapshot."""
        result = self.backend.read(f"crewai:{self.crew_id}:snapshots:{label}")
        if result is None:
            return {"restored": False, "reason": "snapshot_not_found"}

        data = result.get("data", {})
        val = data.get("value", data)
        kb = val.get("knowledge_base", {}) if isinstance(val, dict) else {}

        restored = 0
        for finding in kb.get("findings", []):
            self.store_finding("restored", finding.get("key", ""), finding.get("data", {}))
            restored += 1

        return {"restored": True, "label": label, "items_restored": restored}
