"""
Aniate Engram SDK - Core Client
The middleware pattern for giving your LLM persistent memory.
"""

import json
from typing import Optional, List, Dict, Any, Callable, Union
import requests

from .types import (
    GraphData, QueryResult, StoreResult, IngestResult, 
    ChatResponse, UsageInfo, Entity, Fact
)

class EngramAPIError(Exception):
    """Exception raised for API errors."""
    def __init__(self, message: str, status_code: int, error_code: Optional[str] = None):
        self.message = message
        self.status_code = status_code
        self.error_code = error_code
        super().__init__(f"[{status_code}] {message}")

class Memory:
    """A persistent memory workspace."""
    def __init__(self, client: "Engram", name: str, uid: Optional[str] = None):
        self._client = client
        self.name = name
        self.uid = uid
    
    def store(self, text: str) -> StoreResult:
        """Store raw text in memory (ML extraction)."""
        response = self._client._request(
            "POST",
            f"/memory/{self.name}/store",
            json={"text": text}
        )
        return StoreResult.from_dict(response)
    
    def ingest(
        self, 
        entities: Optional[List[Union[Dict[str, str], Entity]]] = None,
        facts: Optional[List[Union[Dict[str, str], Fact]]] = None
    ) -> IngestResult:
        """Ingest pre-structured entities and facts directly."""
        entity_dicts = []
        if entities:
            for e in entities:
                if isinstance(e, Entity):
                    entity_dicts.append(e.to_dict())
                else:
                    entity_dicts.append(e)
        
        fact_dicts = []
        if facts:
            for f in facts:
                if isinstance(f, Fact):
                    fact_dicts.append(f.to_dict())
                else:
                    fact_dicts.append(f)
        
        response = self._client._request(
            "POST",
            f"/memory/{self.name}/ingest",
            json={"entities": entity_dicts, "facts": fact_dicts}
        )
        return IngestResult.from_dict(response)
    
    def query(self, text: str, top_k: int = 5) -> QueryResult:
        """Query memory using semantic search."""
        response = self._client._request(
            "POST",
            f"/memory/{self.name}/query",
            json={"text": text, "top_k": top_k}
        )
        return QueryResult.from_dict(response)
    
    def graph(self) -> GraphData:
        """Get the full knowledge graph."""
        response = self._client._request("GET", f"/memory/{self.name}/graph")
        return GraphData.from_dict(response)

class Session:
    """An ephemeral chat session."""
    def __init__(self, client: "Engram", uid: str, name: str, base_memory: Optional[str] = None):
        self._client = client
        self.uid = uid
        self.name = name
        self.base_memory = base_memory
    
    def chat(self, message: str) -> ChatResponse:
        """Send a message in this session."""
        response = self._client._request(
            "POST",
            f"/session/{self.uid}/chat",
            json={"message": message}
        )
        return ChatResponse.from_dict(response)
    
    def graph(self) -> GraphData:
        """Get the session's current knowledge graph."""
        response = self._client._request("GET", f"/session/{self.uid}/graph")
        return GraphData.from_dict(response)

class Engram:
    """Aniate Engram Client - Give your LLM persistent memory."""
    def __init__(self, api_key: str, base_url: str = "https://api.aniate.com/v1"):
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self._session = requests.Session()
        self._session.headers.update({
            "x-api-key": api_key,
            "Content-Type": "application/json"
        })
    
    def _request(self, method: str, path: str, json: Optional[Dict] = None) -> Dict[str, Any]:
        url = f"{self.base_url}{path}"
        response = self._session.request(method=method, url=url, json=json)
        
        if not response.ok:
            error_data = response.json() if response.text else {}
            error_msg = error_data.get("message", response.text or response.reason)
            raise EngramAPIError(
                message=error_msg,
                status_code=response.status_code,
                error_code=error_data.get("code")
            )
        return response.json()
    
    def memory(self, name: str) -> Memory:
        try:
            response = self._request("POST", "/memory", json={"memory_name": name})
            return Memory(self, name, response.get("memory_uid"))
        except EngramAPIError as e:
            if e.error_code == "MEMORY_EXISTS":
                return Memory(self, name)
            raise
    
    def session(self, memory_name: Optional[str] = None, session_name: Optional[str] = None) -> Session:
        response = self._request(
            "POST", 
            "/session", 
            json={"memory_name": memory_name, "session_name": session_name}
        )
        return Session(
            self,
            uid=response["session_uid"],
            name=response["session_name"],
            base_memory=response.get("base_memory")
        )
    
    def usage(self) -> UsageInfo:
        response = self._request("GET", "/usage")
        return UsageInfo.from_dict(response)
    
    def list_memories(self) -> List[Dict[str, Any]]:
        response = self._request("GET", "/memory")
        return response.get("memories", [])
