"""
Synrix Agent Runtime — Developer API
The developer-facing interface. Clean, fast, intuitive.

Usage:
    from synrix_runtime import AgentRuntime
    agent = AgentRuntime("my_agent", agent_type="researcher")
    agent.remember("user_preference", {"format": "bullet_points"})
    value = agent.recall("user_preference")
"""

import time
import json
import threading
from typing import Any, List, Optional
from dataclasses import dataclass, field


@dataclass
class MemoryResult:
    node_id: Optional[int]
    key: str
    latency_us: float
    timestamp: float
    success: bool = True


@dataclass
class RecallResult:
    value: Any
    key: str
    latency_us: float
    found: bool


@dataclass
class SearchResult:
    items: list
    count: int
    latency_us: float


@dataclass
class SnapshotResult:
    label: str
    keys_captured: int
    latency_us: float
    size_bytes: int = 0


@dataclass
class RestoreResult:
    label: str
    keys_restored: int
    recovery_time_us: float


@dataclass
class HandoffResult:
    task_id: str
    latency_us: float
    success: bool = True


@dataclass
class TaskResult:
    task_id: str
    payload: dict
    latency_us: float
    found: bool = True


@dataclass
class AgentStats:
    agent_id: str
    total_operations: int = 0
    total_writes: int = 0
    total_reads: int = 0
    total_queries: int = 0
    avg_write_latency_us: float = 0.0
    avg_read_latency_us: float = 0.0
    crash_count: int = 0
    memory_node_count: int = 0
    performance_score: float = 100.0
    uptime_seconds: float = 0.0


class AgentRuntime:
    """
    The Synrix Agent Runtime.

    Usage:
        from synrix_runtime import AgentRuntime
        agent = AgentRuntime("my_agent", agent_type="researcher")
        agent.remember("user_preference", {"format": "bullet_points"})
        value = agent.recall("user_preference")
    """

    def __init__(self, agent_id: str, agent_type: str = "generic", metadata: dict = None):
        self.agent_id = agent_id
        self.agent_type = agent_type
        self.metadata = metadata or {}
        self._write_count = 0
        self._read_count = 0
        self._query_count = 0
        self._started_at = time.time()

        # Connect to Synrix — share daemon's backend for shared memory
        start = time.perf_counter_ns()
        try:
            from synrix_runtime.core.daemon import RuntimeDaemon
            self._daemon = RuntimeDaemon.get_instance()
            if not self._daemon.running:
                self._daemon.start()
            # Reuse the daemon's backend so all agents share the same store
            self.backend = self._daemon.backend
            self._daemon.register_agent(agent_id, agent_type, self.metadata)
        except Exception:
            from synrix.agent_backend import get_synrix_backend
            from synrix_runtime.config import SynrixConfig
            config = SynrixConfig.from_env()
            self.backend = get_synrix_backend(**config.get_backend_kwargs())
            self._daemon = None
        connect_us = (time.perf_counter_ns() - start) / 1000

        # Initialize metrics collector
        try:
            from synrix_runtime.monitoring.metrics import MetricsCollector
            self._metrics = MetricsCollector.get_instance(self.backend)
        except Exception:
            self._metrics = None

        # Start heartbeat thread
        self._heartbeat_running = True
        self._heartbeat_thread = threading.Thread(
            target=self._heartbeat_loop, name=f"heartbeat-{agent_id}", daemon=True
        )
        self._heartbeat_thread.start()

        print(f"[{agent_id}] Runtime connected in {connect_us:.1f}us | type={agent_type}")

    def remember(self, key: str, value: Any, tags: list = None) -> MemoryResult:
        """Store a memory in Synrix."""
        full_key = f"agents:{self.agent_id}:{key}"
        payload = value if isinstance(value, dict) else {"value": value}
        if tags:
            payload["_tags"] = tags

        start = time.perf_counter_ns()
        node_id = self.backend.write(full_key, payload, metadata={"type": "agent_memory", "agent_id": self.agent_id})
        latency_us = (time.perf_counter_ns() - start) / 1000

        self._write_count += 1
        if self._metrics:
            self._metrics.record_write(self.agent_id, key, latency_us, node_id is not None, node_id)

        return MemoryResult(
            node_id=node_id,
            key=key,
            latency_us=latency_us,
            timestamp=time.time(),
            success=node_id is not None,
        )

    def recall(self, key: str) -> RecallResult:
        """Recall a memory from Synrix."""
        full_key = f"agents:{self.agent_id}:{key}"

        start = time.perf_counter_ns()
        result = self.backend.read(full_key)
        latency_us = (time.perf_counter_ns() - start) / 1000

        found = result is not None
        value = None
        if result:
            data = result.get("data", {})
            value = data.get("value", data)

        self._read_count += 1
        if self._metrics:
            self._metrics.record_read(self.agent_id, key, latency_us, found)

        return RecallResult(value=value, key=key, latency_us=latency_us, found=found)

    def search(self, prefix: str, limit: int = 50) -> SearchResult:
        """Search agent memory by prefix."""
        full_prefix = f"agents:{self.agent_id}:{prefix}" if prefix else f"agents:{self.agent_id}:"

        start = time.perf_counter_ns()
        results = self.backend.query_prefix(full_prefix, limit=limit)
        latency_us = (time.perf_counter_ns() - start) / 1000

        items = []
        for r in results:
            key = r.get("key", "")
            # Strip the agent prefix for cleaner results
            short_key = key.replace(f"agents:{self.agent_id}:", "", 1)
            data = r.get("data", {})
            value = data.get("value", data)
            items.append({"key": short_key, "value": value, "node_id": r.get("id")})

        self._query_count += 1
        if self._metrics:
            self._metrics.record_query(self.agent_id, prefix, latency_us, len(items))

        return SearchResult(items=items, count=len(items), latency_us=latency_us)

    def snapshot(self, label: str = None) -> SnapshotResult:
        """Take a snapshot of all current memory."""
        if label is None:
            label = f"snap_{int(time.time()*1000000)}"

        start = time.perf_counter_ns()

        # Get all current memory keys
        all_keys = self.backend.query_prefix(f"agents:{self.agent_id}:", limit=500)
        snapshot_data = {}
        for item in all_keys:
            key = item.get("key", "")
            if ":snapshots:" not in key:
                data = item.get("data", {})
                snapshot_data[key] = data.get("value", data)

        # Write snapshot
        snapshot_payload = {
            "label": label,
            "agent_id": self.agent_id,
            "keys": snapshot_data,
            "key_count": len(snapshot_data),
            "created_at": time.time(),
        }
        size_bytes = len(json.dumps(snapshot_payload).encode())
        self.backend.write(
            f"agents:{self.agent_id}:snapshots:{label}",
            snapshot_payload,
            metadata={"type": "snapshot", "agent_id": self.agent_id}
        )
        latency_us = (time.perf_counter_ns() - start) / 1000

        if self._metrics:
            self._metrics.record_snapshot(self.agent_id, label, len(snapshot_data), latency_us)

        return SnapshotResult(
            label=label,
            keys_captured=len(snapshot_data),
            latency_us=latency_us,
            size_bytes=size_bytes,
        )

    def restore(self, label: str = None) -> RestoreResult:
        """Restore from a named snapshot or the latest one."""
        start = time.perf_counter_ns()

        if label:
            result = self.backend.read(f"agents:{self.agent_id}:snapshots:{label}")
        else:
            # Find latest snapshot
            snapshots = self.backend.query_prefix(f"agents:{self.agent_id}:snapshots:", limit=50)
            if not snapshots:
                return RestoreResult(label="none", keys_restored=0, recovery_time_us=0)
            snapshots.sort(key=lambda x: x.get("data", {}).get("value", {}).get("created_at", 0)
                          if isinstance(x.get("data", {}).get("value"), dict)
                          else x.get("data", {}).get("timestamp", 0), reverse=True)
            result = snapshots[0]
            data = result.get("data", {})
            val = data.get("value", data)
            label = val.get("label", "unknown") if isinstance(val, dict) else "unknown"

        if result is None:
            return RestoreResult(label=label or "none", keys_restored=0, recovery_time_us=0)

        # Extract snapshot data
        data = result.get("data", {})
        val = data.get("value", data)
        keys_data = val.get("keys", {}) if isinstance(val, dict) else {}

        # Restore each key
        restored = 0
        for key, value in keys_data.items():
            self.backend.write(key, value, metadata={"type": "restored", "agent_id": self.agent_id})
            restored += 1

        recovery_us = (time.perf_counter_ns() - start) / 1000

        if self._metrics:
            self._metrics.record_recovery(self.agent_id, recovery_us, restored)

        return RestoreResult(label=label, keys_restored=restored, recovery_time_us=recovery_us)

    def share(self, key: str, value: Any, space: str = "global") -> MemoryResult:
        """Write to shared memory space."""
        full_key = f"shared:{space}:{key}"
        payload = value if isinstance(value, dict) else {"value": value}
        payload["_author"] = self.agent_id
        payload["_shared_at"] = time.time()

        start = time.perf_counter_ns()
        node_id = self.backend.write(full_key, payload, metadata={"type": "shared_memory", "space": space, "author": self.agent_id})
        latency_us = (time.perf_counter_ns() - start) / 1000

        # Write changelog
        ts = int(time.time() * 1000000)
        self.backend.write(
            f"shared:{space}:changelog:{ts}",
            {"key": key, "author": self.agent_id, "action": "write", "timestamp": time.time()},
            metadata={"type": "shared_changelog"}
        )

        self._write_count += 1
        if self._metrics:
            self._metrics.record_write(self.agent_id, f"shared:{space}:{key}", latency_us, True, node_id)

        return MemoryResult(node_id=node_id, key=key, latency_us=latency_us, timestamp=time.time())

    def read_shared(self, key: str, space: str = "global") -> RecallResult:
        """Read from shared memory space."""
        full_key = f"shared:{space}:{key}"

        start = time.perf_counter_ns()
        result = self.backend.read(full_key)
        latency_us = (time.perf_counter_ns() - start) / 1000

        found = result is not None
        value = None
        if result:
            data = result.get("data", {})
            value = data.get("value", data)

        self._read_count += 1
        if self._metrics:
            self._metrics.record_read(self.agent_id, f"shared:{space}:{key}", latency_us, found)

        return RecallResult(value=value, key=key, latency_us=latency_us, found=found)

    def subscribe_shared(self, space: str, callback: callable):
        """Poll a shared space for new keys and call callback on changes."""
        seen_keys = set()

        def _poll():
            while self._heartbeat_running:
                try:
                    results = self.backend.query_prefix(f"shared:{space}:", limit=200)
                    for r in results:
                        key = r.get("key", "")
                        if key not in seen_keys:
                            seen_keys.add(key)
                            data = r.get("data", {})
                            value = data.get("value", data)
                            try:
                                callback(key, value)
                            except Exception:
                                pass
                except Exception:
                    pass
                time.sleep(1)

        t = threading.Thread(target=_poll, name=f"sub-{space}-{self.agent_id}", daemon=True)
        t.start()

    def handoff(self, task_id: str, to_agent: str, payload: dict) -> HandoffResult:
        """Hand off a task to another agent."""
        handoff_data = {
            "task_id": task_id,
            "from_agent": self.agent_id,
            "to_agent": to_agent,
            "payload": payload,
            "status": "pending",
            "created_at": time.time(),
        }

        start = time.perf_counter_ns()
        self.backend.write(
            f"tasks:handoff:{task_id}",
            handoff_data,
            metadata={"type": "task_handoff", "from": self.agent_id, "to": to_agent}
        )
        latency_us = (time.perf_counter_ns() - start) / 1000

        # Audit event
        try:
            from synrix_runtime.monitoring.audit import AuditSystem
            audit = AuditSystem(self.backend)
            audit.log_handoff(self.agent_id, to_agent, task_id, payload, {})
        except Exception:
            pass

        if self._metrics:
            self._metrics.record_handoff(self.agent_id, to_agent, task_id, latency_us)

        return HandoffResult(task_id=task_id, latency_us=latency_us)

    def claim_task(self, task_id: str) -> TaskResult:
        """Claim a pending task."""
        start = time.perf_counter_ns()
        result = self.backend.read(f"tasks:handoff:{task_id}")
        latency_us = (time.perf_counter_ns() - start) / 1000

        if result is None:
            return TaskResult(task_id=task_id, payload={}, latency_us=latency_us, found=False)

        data = result.get("data", {})
        val = data.get("value", data)

        # Update status
        if isinstance(val, dict):
            val["status"] = "claimed"
            val["claimed_by"] = self.agent_id
            val["claimed_at"] = time.time()
        self.backend.write(f"tasks:handoff:{task_id}", val, metadata={"type": "task_claimed"})

        return TaskResult(task_id=task_id, payload=val, latency_us=latency_us)

    def complete_task(self, task_id: str, result: dict) -> TaskResult:
        """Mark a task as complete with result."""
        completion = {
            "task_id": task_id,
            "completed_by": self.agent_id,
            "result": result,
            "completed_at": time.time(),
        }

        start = time.perf_counter_ns()
        self.backend.write(f"tasks:complete:{task_id}", completion, metadata={"type": "task_complete"})
        latency_us = (time.perf_counter_ns() - start) / 1000

        # Audit
        try:
            from synrix_runtime.monitoring.audit import AuditSystem
            audit = AuditSystem(self.backend)
            audit.log_handoff(self.agent_id, "system", task_id, completion, {})
        except Exception:
            pass

        return TaskResult(task_id=task_id, payload=completion, latency_us=latency_us)

    def log_decision(self, decision: str, reasoning: str, context: dict = None):
        """Log a decision with full audit trail."""
        ts = int(time.time() * 1000000)

        # Capture memory snapshot at decision time
        memory_snapshot = {}
        try:
            all_keys = self.backend.query_prefix(f"agents:{self.agent_id}:", limit=200)
            for item in all_keys:
                key = item.get("key", "")
                if ":snapshots:" not in key:
                    data = item.get("data", {})
                    memory_snapshot[key] = data.get("value", data)
        except Exception:
            pass

        decision_data = {
            "agent_id": self.agent_id,
            "decision": decision,
            "reasoning": reasoning,
            "context": context or {},
            "memory_snapshot": memory_snapshot,
            "timestamp": time.time(),
        }

        self.backend.write(
            f"audit:{self.agent_id}:{ts}:decision",
            decision_data,
            metadata={"type": "audit_decision", "agent_id": self.agent_id}
        )

    def get_stats(self) -> AgentStats:
        """Get complete performance statistics."""
        if self._metrics:
            m = self._metrics.get_agent_metrics(self.agent_id)
            return AgentStats(
                agent_id=self.agent_id,
                total_operations=m.total_operations,
                total_writes=m.total_writes,
                total_reads=m.total_reads,
                total_queries=m.total_queries,
                avg_write_latency_us=m.avg_write_latency_us,
                avg_read_latency_us=m.avg_read_latency_us,
                crash_count=m.crash_count,
                memory_node_count=m.memory_node_count,
                performance_score=m.performance_score,
                uptime_seconds=time.time() - self._started_at,
            )
        return AgentStats(
            agent_id=self.agent_id,
            total_operations=self._write_count + self._read_count + self._query_count,
            total_writes=self._write_count,
            total_reads=self._read_count,
            total_queries=self._query_count,
            uptime_seconds=time.time() - self._started_at,
        )

    def _heartbeat_loop(self):
        """Background heartbeat — writes to Synrix every 5 seconds."""
        while self._heartbeat_running:
            try:
                now = time.time()
                self.backend.write(
                    f"runtime:agents:{self.agent_id}:heartbeat",
                    {"value": now},
                    metadata={"type": "heartbeat"}
                )
                self.backend.write(
                    f"runtime:agents:{self.agent_id}:last_active",
                    {"value": now},
                    metadata={"type": "timestamp"}
                )
            except Exception:
                pass
            time.sleep(5)

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.shutdown()

    def shutdown(self):
        """Gracefully shut down the agent."""
        self._heartbeat_running = False

        # Final snapshot
        try:
            self.snapshot("shutdown_auto")
        except Exception:
            pass

        # Deregister
        try:
            if self._daemon:
                self._daemon.deregister_agent(self.agent_id)
        except Exception:
            pass

        print(f"[{self.agent_id}] Shutdown complete | writes={self._write_count} reads={self._read_count}")
