from __future__ import annotations

import traceback
from dataclasses import dataclass, field
from typing import List, Literal, Optional, Type, Union

from pydantic import BaseModel
from openai.types.chat import ChatCompletionMessageParam

from mito_ai.prompt_builders.chat_prompt import create_chat_prompt
from mito_ai.prompt_builders.inline_completer_prompt import create_inline_prompt
from mito_ai.prompt_builders.explain_code_prompt import create_explain_code_prompt
from mito_ai.prompt_builders.smart_debug_prompt import create_error_prompt
from mito_ai.prompt_builders.agent_planning_prompt import create_agent_prompt

CompletionIncomingMessageTypes = Literal['chat', 'inline_completion', 'codeExplain', 'smartDebug', 'agent:planning']
IncomingMessageTypes = Union[Literal['clear_history', 'fetch_history'], CompletionIncomingMessageTypes]

@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 ChatMessageBuilder:
    variables: Optional[List[str]] = None
    activeCellCode: Optional[str] = None
    input: Optional[str] = None
    index: Optional[int] = None

    @property
    def prompt(self) -> str:
        return create_chat_prompt(self.variables or [], self.activeCellCode or '', self.input or '')
    
    @property
    def display_message(self) -> str:
        cell_code_block = f"""```python
{self.activeCellCode}
```

"""
        return f"{cell_code_block if self.activeCellCode else ''}{self.input}"
    
    @property
    def pro_model(self) -> str:
        return "o3-mini"
    
    @property
    def os_model(self) -> str:
        return "gpt-4o-mini"

@dataclass(frozen=True)
class SmartDebugMessageBuilder:
    variables: Optional[List[str]] = None
    activeCellCode: Optional[str] = None
    errorMessage: Optional[str] = None

    @property
    def prompt(self) -> str:
        return create_error_prompt(self.errorMessage or '', self.activeCellCode or '', self.variables or [])
    
    @property
    def display_message(self) -> str:
        cell_code_block = f"""```python
{self.activeCellCode}
```

"""
        return f"{cell_code_block if self.activeCellCode else ''}{self.errorMessage}"
    
    @property
    def pro_model(self) -> str:
        return "gpt-4o-mini"
    
    @property
    def os_model(self) -> str:
        return "gpt-4o-mini"

@dataclass(frozen=True)
class CodeExplainMessageBuilder:
    variables: Optional[List[str]] = None
    activeCellCode: Optional[str] = None

    @property
    def prompt(self) -> str:
        return create_explain_code_prompt(self.activeCellCode or '')
    
    @property
    def display_message(self) -> str:
        cell_code_block = f"""```python
{self.activeCellCode}
```

"""
        
        return f"{cell_code_block if self.activeCellCode else ''}Explain this code"
    
    @property
    def pro_model(self) -> str:
        return "gpt-4o-mini"
    
    @property
    def os_model(self) -> str:
        return "gpt-4o-mini"

@dataclass(frozen=True)
class InlineCompletionMessageBuilder:
    prefix: Optional[str] = None
    suffix: Optional[str] = None
    variables: Optional[List[str]] = None

    @property
    def prompt(self) -> str:
        return create_inline_prompt(self.prefix or '', self.suffix or '', self.variables or [])
    
    @property
    def pro_model(self) -> str:
        return "gpt-4o-mini"
    
    @property
    def os_model(self) -> str:
        return "gpt-4o-mini"

@dataclass(frozen=True)
class AgentMessageBuilder:
    fileType: Optional[str] = None
    columnSamples: Optional[List[str]] = None
    input: Optional[str] = None
    variables: Optional[List[str]] = None

    @property
    def prompt(self) -> str:
        return create_agent_prompt(
            self.fileType or "",
            self.columnSamples or [],
            self.input or "",
            self.variables or [],
        )

    @property
    def display_message(self) -> str:
        return self.input or ''
    
    @property
    def pro_model(self) -> str:
        return "o3-mini"
    
    @property
    def os_model(self) -> str:
        return "gpt-4o-mini"

    @property
    def response_format(self) -> Type[BaseModel]:
        class PlanOfAttack(BaseModel):
            actions: List[str]
            dependencies: List[str]
        return PlanOfAttack
    


@dataclass(frozen=True)
class CompletionRequest:
    """
    Message send by the client to request an AI chat response.
    """

    # Message type.
    type: IncomingMessageTypes

    # 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 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"

