"""
Type definitions for Aniate Engram SDK
"""

from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any

@dataclass
class Entity:
    """An extracted entity (node in the knowledge graph)"""
    name: str
    type: str
    
    def to_dict(self) -> Dict[str, str]:
        return {"name": self.name, "type": self.type}

@dataclass
class Fact:
    """A relationship between entities (edge in the knowledge graph)"""
    subject: str
    predicate: str
    object: str
    
    def to_dict(self) -> Dict[str, str]:
        return {
            "subject": self.subject,
            "predicate": self.predicate,
            "object": self.object
        }

@dataclass
class GraphNode:
    """A node in the knowledge graph"""
    id: str
    label: str
    type: str
    embedding: Optional[List[float]] = None

@dataclass
class GraphEdge:
    """An edge in the knowledge graph"""
    source: str
    target: str
    label: str

@dataclass
class GraphData:
    """Complete knowledge graph data (D3-compatible)"""
    nodes: List[GraphNode] = field(default_factory=list)
    edges: List[GraphEdge] = field(default_factory=list)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "GraphData":
        nodes = [
            GraphNode(
                id=n.get("id", ""),
                label=n.get("label", ""),
                type=n.get("type", ""),
                embedding=n.get("embedding")
            )
            for n in data.get("nodes", [])
        ]
        edges = [
            GraphEdge(
                source=e.get("source", ""),
                target=e.get("target", ""),
                label=e.get("label", "")
            )
            for e in data.get("edges", [])
        ]
        return cls(nodes=nodes, edges=edges)

@dataclass
class SearchResult:
    """A single search result from semantic query"""
    id: str
    text: str
    entity_type: str
    score: float

@dataclass
class QueryResult:
    """Result from a memory query"""
    results: List[SearchResult] = field(default_factory=list)
    graph: Optional[GraphData] = None
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "QueryResult":
        results = [
            SearchResult(
                id=r.get("id", ""),
                text=r.get("text", ""),
                entity_type=r.get("entity_type", ""),
                score=r.get("score", 0.0)
            )
            for r in data.get("results", [])
        ]
        graph = GraphData.from_dict(data.get("graph", {})) if data.get("graph") else None
        return cls(results=results, graph=graph)

@dataclass
class StoreResult:
    """Result from storing data in memory"""
    nodes_added: int
    edges_added: int
    graph: GraphData
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "StoreResult":
        return cls(
            nodes_added=data.get("nodes_added", 0),
            edges_added=data.get("edges_added", 0),
            graph=GraphData.from_dict(data.get("graph", {}))
        )

@dataclass
class IngestResult:
    """Result from ingesting structured data"""
    nodes_added: int
    edges_added: int
    facts_added: int
    graph: GraphData
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "IngestResult":
        return cls(
            nodes_added=data.get("nodes_added", 0),
            edges_added=data.get("edges_added", 0),
            facts_added=data.get("facts_added", 0),
            graph=GraphData.from_dict(data.get("graph", {}))
        )

@dataclass
class ChatResponse:
    """Response from a chat session"""
    response: str
    entities: List[Dict[str, Any]]
    graph: GraphData
    usage: Dict[str, Any]
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ChatResponse":
        return cls(
            response=data.get("response", ""),
            entities=data.get("entities", []),
            graph=GraphData.from_dict(data.get("graph", {})),
            usage=data.get("usage", {})
        )

@dataclass
class UsageInfo:
    """User's API usage information"""
    queries_used: int
    queries_limit: int
    plan: str
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "UsageInfo":
        return cls(
            queries_used=data.get("queries_used", 0),
            queries_limit=data.get("queries_limit", 1000),
            plan=data.get("plan", "hobby")
        )
