"""
Kronvex Python SDK v0.3.0
Persistent memory layer for AI agents.

Usage:
    from kronvex import Kronvex

    kx = Kronvex(api_key="kx_live_...", agent_id="<uuid>")
    kx.remember("User prefers dark mode", memory_type="semantic")
    results = kx.recall("user preferences")
    context = kx.inject_context("What theme should I use?")
"""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Optional
from uuid import UUID

try:
    import httpx
    _HTTPX = True
except ImportError:
    import urllib.request
    import json as _json
    _HTTPX = False


BASE_URL = "https://api.kronvex.io/api/v1"


# ── Data classes ──────────────────────────────────────────────────────────────

@dataclass
class Memory:
    id: str
    agent_id: str
    content: str
    memory_type: str
    session_id: Optional[str]
    created_at: str
    last_accessed_at: Optional[str]
    access_count: int
    expires_at: Optional[str]
    pinned: bool
    metadata: dict

    @classmethod
    def _from_dict(cls, d: dict) -> "Memory":
        return cls(
            id=d["id"],
            agent_id=d["agent_id"],
            content=d["content"],
            memory_type=d.get("memory_type", "episodic"),
            session_id=d.get("session_id"),
            created_at=d.get("created_at", ""),
            last_accessed_at=d.get("last_accessed_at"),
            access_count=d.get("access_count", 0),
            expires_at=d.get("expires_at"),
            pinned=d.get("pinned", False),
            metadata=d.get("metadata", {}),
        )


@dataclass
class RecallResult:
    memory: Memory
    similarity: float
    confidence: float


@dataclass
class RecallResponse:
    query: str
    results: list[RecallResult]
    total_found: int

    def __iter__(self):
        return iter(self.results)

    def best(self) -> Optional[RecallResult]:
        return self.results[0] if self.results else None


@dataclass
class InjectContextResponse:
    context_block: str
    memories_used: int
    memories: list[RecallResult]

    def __str__(self) -> str:
        return self.context_block


# ── Exceptions ────────────────────────────────────────────────────────────────

class KronvexError(Exception):
    def __init__(self, message: str, status_code: int = 0):
        super().__init__(message)
        self.status_code = status_code


class AuthError(KronvexError): pass
class RateLimitError(KronvexError): pass
class NotFoundError(KronvexError): pass


# ── Client ────────────────────────────────────────────────────────────────────

class Kronvex:
    """
    Kronvex API client.

    Args:
        api_key:  Your Kronvex API key (kx_live_...). Reads KRONVEX_API_KEY env var if not provided.
        agent_id: Default agent UUID. Can be overridden per call.
        base_url: Override API base URL (useful for self-hosted).
        timeout:  Request timeout in seconds (default: 10).
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        agent_id: Optional[str] = None,
        base_url: str = BASE_URL,
        timeout: float = 10.0,
    ):
        self.api_key = api_key or os.getenv("KRONVEX_API_KEY")
        if not self.api_key:
            raise AuthError("No API key provided. Pass api_key= or set KRONVEX_API_KEY env var.")
        self.agent_id = agent_id or os.getenv("KRONVEX_AGENT_ID")
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self._headers = {
            "X-API-Key": self.api_key,
            "Content-Type": "application/json",
            "User-Agent": "kronvex-python/0.3.0",
        }

    # ── Core methods ──────────────────────────────────────────────────────────

    def remember(
        self,
        content: str,
        *,
        memory_type: str = "episodic",
        session_id: Optional[str] = None,
        metadata: Optional[dict] = None,
        ttl_days: Optional[int] = None,
        pinned: bool = False,
        agent_id: Optional[str] = None,
    ) -> Memory:
        """
        Store a memory for an agent.

        Args:
            content:     The text to remember.
            memory_type: "episodic" | "semantic" | "procedural" (default: "episodic").
            session_id:  Optional session identifier for grouping memories.
            metadata:    Arbitrary JSON metadata dict.
            ttl_days:    Days until memory expires (1–3650). None = never expires.
            pinned:      Pinned memories never expire, even with ttl_days set.
            agent_id:    Override the default agent_id.

        Returns:
            Memory object with id, content, expires_at, etc.

        Example:
            mem = kx.remember(
                "User is a power user who prefers keyboard shortcuts",
                memory_type="semantic",
                session_id="session_abc",
                ttl_days=90,
            )
        """
        payload: dict[str, Any] = {
            "content": content,
            "memory_type": memory_type,
        }
        if session_id:  payload["session_id"] = session_id
        if metadata:    payload["metadata"] = metadata
        if ttl_days:    payload["ttl_days"] = ttl_days
        if pinned:      payload["pinned"] = True

        aid = self._resolve_agent(agent_id)
        data = self._post(f"/agents/{aid}/remember", payload)
        return Memory._from_dict(data)

    def recall(
        self,
        query: str,
        *,
        top_k: int = 5,
        threshold: float = 0.5,
        session_id: Optional[str] = None,
        memory_type: Optional[str] = None,
        agent_id: Optional[str] = None,
    ) -> RecallResponse:
        """
        Recall relevant memories by semantic similarity.

        Args:
            query:       Natural language query.
            top_k:       Max number of results (default: 5).
            threshold:   Minimum similarity score 0–1 (default: 0.5).
            session_id:  Filter to a specific session.
            memory_type: Filter by type ("episodic" | "semantic" | "procedural").
            agent_id:    Override the default agent_id.

        Returns:
            RecallResponse with .results (list of RecallResult with .memory, .similarity, .confidence).

        Example:
            resp = kx.recall("user payment preferences", top_k=3)
            for r in resp:
                print(r.confidence, r.memory.content)
        """
        payload: dict[str, Any] = {"query": query, "top_k": top_k, "threshold": threshold}
        if session_id:   payload["session_id"] = session_id
        if memory_type:  payload["memory_type"] = memory_type

        aid = self._resolve_agent(agent_id)
        data = self._post(f"/agents/{aid}/recall", payload)
        results = [
            RecallResult(
                memory=Memory._from_dict(r["memory"]),
                similarity=r.get("similarity", 0.0),
                confidence=r.get("confidence", 0.0),
            )
            for r in data.get("results", [])
        ]
        return RecallResponse(
            query=data.get("query", query),
            results=results,
            total_found=data.get("total_found", len(results)),
        )

    def inject_context(
        self,
        message: str,
        *,
        top_k: int = 5,
        threshold: float = 0.5,
        agent_id: Optional[str] = None,
    ) -> InjectContextResponse:
        """
        Inject relevant memory context for an incoming user message.
        Use the returned .context_block as a prefix to your system prompt.

        Args:
            message:   The current user message.
            top_k:     Max memories to inject (default: 5).
            threshold: Minimum similarity threshold (default: 0.5).
            agent_id:  Override the default agent_id.

        Returns:
            InjectContextResponse. Cast to str or use .context_block directly.

        Example:
            ctx = kx.inject_context(user_message)
            system_prompt = str(ctx) + "\\n\\nYou are a helpful assistant."
        """
        payload: dict[str, Any] = {"message": message, "top_k": top_k, "threshold": threshold}
        aid = self._resolve_agent(agent_id)
        data = self._post(f"/agents/{aid}/inject-context", payload)
        results = [
            RecallResult(
                memory=Memory._from_dict(r["memory"]),
                similarity=r.get("similarity", 0.0),
                confidence=r.get("confidence", 0.0),
            )
            for r in data.get("memories", [])
        ]
        return InjectContextResponse(
            context_block=data.get("context_block", ""),
            memories_used=data.get("memories_used", 0),
            memories=results,
        )

    # ── Agent helpers ─────────────────────────────────────────────────────────

    def create_agent(self, name: str, description: str = "") -> dict:
        """Create a new agent and return its data dict (id, name, created_at)."""
        return self._post("/agents", {"name": name, "description": description})

    def list_agents(self) -> list[dict]:
        """List all agents for this API key."""
        return self._get("/agents")

    def list_sessions(self, agent_id: Optional[str] = None) -> list[dict]:
        """List distinct session_ids for an agent with memory counts."""
        aid = self._resolve_agent(agent_id)
        data = self._get(f"/agents/{aid}/sessions")
        return data.get("sessions", [])

    def delete_memory(self, memory_id: str, agent_id: Optional[str] = None) -> bool:
        """Delete a specific memory by ID. Returns True if deleted."""
        aid = self._resolve_agent(agent_id)
        self._delete(f"/agents/{aid}/memories/{memory_id}")
        return True

    # ── Context manager support ───────────────────────────────────────────────

    def __enter__(self): return self
    def __exit__(self, *_): pass

    # ── HTTP helpers ──────────────────────────────────────────────────────────

    def _resolve_agent(self, agent_id: Optional[str]) -> str:
        aid = agent_id or self.agent_id
        if not aid:
            raise ValueError("No agent_id provided. Set it in Kronvex(agent_id=...) or pass per call.")
        return aid

    def _post(self, path: str, payload: dict) -> dict:
        return self._request("POST", path, payload)

    def _get(self, path: str) -> Any:
        return self._request("GET", path)

    def _delete(self, path: str) -> None:
        self._request("DELETE", path)

    def _request(self, method: str, path: str, payload: Optional[dict] = None) -> Any:
        url = self.base_url + path
        if _HTTPX:
            r = httpx.request(
                method, url,
                headers=self._headers,
                json=payload,
                timeout=self.timeout,
            )
            self._raise_for_status(r.status_code, r.text)
            return r.json() if r.content else {}
        else:
            # Stdlib fallback
            import json
            body = json.dumps(payload).encode() if payload else None
            req = urllib.request.Request(url, data=body, headers=self._headers, method=method)
            try:
                with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                    return json.loads(resp.read())
            except urllib.error.HTTPError as e:
                self._raise_for_status(e.code, e.read().decode())

    @staticmethod
    def _raise_for_status(status: int, body: str) -> None:
        if status < 400:
            return
        msg = f"HTTP {status}: {body[:200]}"
        if status == 401: raise AuthError(msg, status)
        if status == 404: raise NotFoundError(msg, status)
        if status == 429: raise RateLimitError(msg, status)
        raise KronvexError(msg, status)
