import traceback
from dataclasses import dataclass, field
from typing import List, Literal, Optional, Type, Union, NewType
from openai.types.chat import ChatCompletionMessageParam
from enum import Enum
from pydantic import BaseModel

# The ThreadID is the unique identifier for the chat thread.
ThreadID = NewType('ThreadID', str)

@dataclass(frozen=True)
class AIOptimizedCell():
  cell_type: str
  id: str
  code: str
  
# Response format for agent planning
# TODO: Figure out how to discriminate the 
# CellUpdateModification and CellUpdateNew types
    
class CellUpdate(BaseModel):
    type: Literal['new', 'modification']
    index: Optional[int]
    id: Optional[str]
    code: str
    description: str
    
class AgentResponse(BaseModel):
    is_finished: bool
    message: str
    cell_update: Optional[CellUpdate]
  
@dataclass(frozen=True)
class ResponseFormatInfo():
    name: str
    # Use the type because we are actually just providing the type format, not an actual instance of the format
    format: type[AgentResponse]


class MessageType(Enum):
    """
    This is all of the different types of messages that we support through the on_message handler.
    """
    CHAT = "chat"
    SMART_DEBUG = "smartDebug"
    CODE_EXPLAIN = "codeExplain"
    AGENT_EXECUTION = "agent:execution"
    AGENT_AUTO_ERROR_FIXUP = "agent:autoErrorFixup"
    INLINE_COMPLETION = "inline_completion"
    CHAT_NAME_GENERATION = "chat_name_generation"
    START_NEW_CHAT = "start_new_chat"
    FETCH_HISTORY = "fetch_history"
    GET_THREADS = "get_threads"
    DELETE_THREAD = "delete_thread"

@dataclass(frozen=True)
class ChatMessageMetadata():
    promptType: Literal['chat']
    input: str
    variables: Optional[List[str]] = None
    files: Optional[List[str]] = None
    activeCellCode: Optional[str] = None
    index: Optional[int] = None   
    
@dataclass(frozen=True)
class AgentExecutionMetadata():
    promptType: Literal['agent:execution']
    input: str
    aiOptimizedCells: List[AIOptimizedCell]
    variables: Optional[List[str]] = None
    files: Optional[List[str]] = None
    
@dataclass(frozen=True)
class AgentSmartDebugMetadata():
    promptType: Literal['agent:autoErrorFixup']
    aiOptimizedCells: List[AIOptimizedCell]
    errorMessage: str
    error_message_producing_code_cell_id: str
    variables: Optional[List[str]] = None
    files: Optional[List[str]] = None
    
@dataclass(frozen=True)
class SmartDebugMetadata():
    promptType: Literal['smartDebug']
    errorMessage: str
    variables: Optional[List[str]] = None
    files: Optional[List[str]] = None
    activeCellCode: Optional[str] = None
    
@dataclass(frozen=True)
class CodeExplainMetadata():    
    promptType: Literal['codeExplain']
    variables: Optional[List[str]] = None
    activeCellCode: Optional[str] = None
    
@dataclass(frozen=True)
class InlineCompleterMetadata():
    promptType: Literal['inline_completion']
    prefix: str 
    suffix: str
    variables: Optional[List[str]] = None
    files: Optional[List[str]] = None

@dataclass(frozen=True)
class FetchHistoryMetadata():
    promptType: Literal['fetch_history']
    
    
@dataclass(frozen=True)
class CompletionRequest:
    """
    Message send by the client to request an AI chat response.
    """

    # Message type.
    type: MessageType

    # Message UID generated by the client.
    message_id: str

    # Chat messages.
    messages: List[ChatCompletionMessageParam] = field(default_factory=list)

    # Whether to stream the response (if supported by the model).
    stream: bool = False
    
    
@dataclass(frozen=True)
class AICapabilities:
    """
    AI provider capabilities
    """

    # Configuration schema.
    configuration: dict

    # AI provider name.
    provider: str

    # Message type.
    type: str = "ai_capabilities"


@dataclass(frozen=True)
class CompletionItemError:
    """
    Completion item error information.
    """

    # Error message.
    message: Optional[str] = None


@dataclass(frozen=True)
class CompletionItem:
    """
    A completion suggestion.
    """

    # The completion.
    content: str

    # Whether the completion is incomplete or not.
    isIncomplete: Optional[bool] = None
    
    # Unique token identifying the completion request in the frontend.
    token: Optional[str] = None

    # Error information for the completion item.
    error: Optional[CompletionItemError] = None


@dataclass(frozen=True)
class CompletionError:
    """
    Completion error description.
    """

    # Error type.
    error_type: str

    # Error title.
    title: str

    # Error traceback.
    traceback: str

    # Hint to resolve the error.
    hint: str = ""

    @staticmethod
    def from_exception(exception: BaseException, hint: str = "") -> "CompletionError":
        """
        Create a completion error from an exception.
        
        Note: OpenAI exceptions can include a 'body' attribute with detailed error information.
        While mypy doesn't know about this attribute on BaseException, we need to handle it
        to properly extract error messages from OpenAI API responses.
        """
        error_type = type(exception)
        error_module = getattr(error_type, "__module__", "")
        return CompletionError(
            error_type=f"{error_module}.{error_type.__name__}"
            if error_module
            else error_type.__name__,
            title=exception.body.get("message")
            if hasattr(exception, "body")
            else (exception.args[0] if exception.args else "Exception"),
            traceback=traceback.format_exc(),
            hint=hint,
        )


@dataclass(frozen=True)
class ErrorMessage(CompletionError):
    """
    Error message.
    """

    # Message type.
    type: Literal["error"] = "error"



@dataclass(frozen=True)
class CompletionReply:
    """
    Message sent from model to client with the completion suggestions.
    """

    # List of completion items.
    items: List[CompletionItem]

    # Parent message UID.
    parent_id: str

    # Message type.
    type: Literal["reply"] = "reply"

    # Completion error.
    error: Optional[CompletionError] = None


@dataclass(frozen=True)
class CompletionStreamChunk:
    """
    Message sent from model to client with the infill suggestions
    """

    chunk: CompletionItem

    # Parent message UID.
    parent_id: str

    # Whether the completion is done or not.
    done: bool

    # Message type.
    type: Literal["chunk"] = "chunk"

    # Completion error.
    error: Optional[CompletionError] = None
    """Completion error."""

@dataclass(frozen=True)
class FetchHistoryReply:
    """
    Message sent from model to client with the chat history.
    """

    # Message UID.
    parent_id: str

    # List of chat messages.
    items: List[ChatCompletionMessageParam]

    # Message type.
    type: Literal["reply"] = "reply"

@dataclass(frozen=True)
class ChatThreadMetadata:
    """
    Chat thread item.
    """

    thread_id: ThreadID

    name: str

    creation_ts: float

    last_interaction_ts: float

@dataclass(frozen=True)
class StartNewChatReply:
    """
    Message sent from model to client after starting a new chat thread.
    """

    # Message UID.
    parent_id: str

    # Chat thread item.
    thread_id: ThreadID

    # Message type.
    type: Literal["reply"] = "reply"

@dataclass(frozen=True)
class FetchThreadsReply:
    """
    Message sent from model to client with the chat threads.
    """

    # Message UID.
    parent_id: str

    # List of chat threads.
    threads: List[ChatThreadMetadata]

    # Message type.
    type: Literal["reply"] = "reply"

@dataclass(frozen=True)
class DeleteThreadReply:
    """
    Message sent from model to client after deleting a chat thread.
    """

    # Message UID.
    parent_id: str

    #Success message
    success: bool

    # Message type.
    type: Literal["reply"] = "reply"