"""Simons AI Assistant response models for the CPZ SDK."""
from __future__ import annotations

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


@dataclass
class TokenUsage:
    """Token usage statistics from an AI response."""
    input_tokens: int = 0
    output_tokens: int = 0
    total_tokens: int = 0
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TokenUsage":
        return cls(
            input_tokens=data.get("input_tokens", data.get("prompt_tokens", 0)),
            output_tokens=data.get("output_tokens", data.get("completion_tokens", 0)),
            total_tokens=data.get("total_tokens", 0),
        )


@dataclass
class CodeProposal:
    """A code change proposal from Simons."""
    file_path: str
    diff: str
    description: str = ""
    applied: bool = False
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "CodeProposal":
        return cls(
            file_path=data.get("file_path", data.get("path", "")),
            diff=data.get("diff", data.get("content", "")),
            description=data.get("description", ""),
            applied=data.get("applied", False),
        )


@dataclass
class SimonsResponse:
    """Complete response from Simons AI Assistant.
    
    Attributes:
        content: The main response text from Simons
        model: The AI model used (e.g., 'claude-sonnet-4-5-20250929')
        tokens_used: Token usage statistics
        cost: Estimated cost in USD for this request
        thinking: Extended thinking content (if show_thinking was enabled)
        proposals: Code change proposals (in agent mode)
        raw: The raw response dictionary from the API
    """
    content: str
    model: str = ""
    tokens_used: TokenUsage = field(default_factory=TokenUsage)
    cost: float = 0.0
    thinking: Optional[str] = None
    proposals: Optional[List[CodeProposal]] = None
    raw: Dict[str, Any] = field(default_factory=dict)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "SimonsResponse":
        # Extract content from various possible response formats
        content = (
            data.get("content") or 
            data.get("response") or 
            data.get("message") or 
            data.get("text") or
            ""
        )
        
        # Parse token usage
        usage_data = data.get("usage", data.get("tokens", {}))
        tokens_used = TokenUsage.from_dict(usage_data) if usage_data else TokenUsage()
        
        # Parse proposals if present
        proposals = None
        proposals_data = data.get("proposals", [])
        if proposals_data:
            proposals = [CodeProposal.from_dict(p) for p in proposals_data]
        
        return cls(
            content=content,
            model=data.get("model", ""),
            tokens_used=tokens_used,
            cost=data.get("cost", 0.0),
            thinking=data.get("thinking"),
            proposals=proposals,
            raw=data,
        )


@dataclass
class SimonsChunk:
    """A streaming chunk from Simons AI Assistant.
    
    Attributes:
        type: The type of chunk ('text', 'thinking', 'done', 'error')
        content: The chunk content
        delta: For text chunks, the incremental text delta
    """
    type: Literal["text", "thinking", "done", "error"]
    content: str
    delta: str = ""
    
    @classmethod
    def from_sse(cls, event_type: str, data: str) -> "SimonsChunk":
        """Parse a Server-Sent Event into a chunk."""
        chunk_type: Literal["text", "thinking", "done", "error"] = "text"
        
        if event_type == "done" or data == "[DONE]":
            chunk_type = "done"
        elif event_type == "error":
            chunk_type = "error"
        elif event_type == "thinking":
            chunk_type = "thinking"
        
        return cls(type=chunk_type, content=data, delta=data if chunk_type == "text" else "")


@dataclass
class SimonsMemory:
    """User's Simons memory containing preferences and learned facts.
    
    Attributes:
        user_id: The user's ID
        preferences: User preferences (risk tolerance, trading style, etc.)
        facts: Learned facts about the user from conversations
        created_at: When the memory was created
        updated_at: When the memory was last updated
    """
    user_id: str = ""
    preferences: Dict[str, Any] = field(default_factory=dict)
    facts: List[str] = field(default_factory=list)
    created_at: Optional[str] = None
    updated_at: Optional[str] = None
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "SimonsMemory":
        return cls(
            user_id=data.get("user_id", ""),
            preferences=data.get("preferences", {}),
            facts=data.get("facts", []),
            created_at=data.get("created_at"),
            updated_at=data.get("updated_at"),
        )
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert memory to dictionary for API calls."""
        return {
            "user_id": self.user_id,
            "preferences": self.preferences,
            "facts": self.facts,
        }
