from __future__ import annotations

import traceback
from enum import Enum
from types import TracebackType

from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import field_serializer

from .core import ErrorResult
from .core import Result
from .core import StatementResult


class DjangoShellOutput(BaseModel):
    status: ExecutionStatus
    output: Output
    stdout: str
    stderr: str

    @classmethod
    def from_result(cls, result: Result) -> DjangoShellOutput:
        output: Output

        match result:
            case StatementResult():
                output = StatementOutput()
            case ErrorResult():
                exception = ExceptionOutput(
                    exc_type=type(result.exception),
                    message=str(result.exception),
                    traceback=result.exception.__traceback__,
                )
                output = ErrorOutput(exception=exception)

        return cls(
            status=ExecutionStatus.from_output(output),
            output=output,
            stdout=result.stdout,
            stderr=result.stderr,
        )


class ExecutionStatus(str, Enum):
    SUCCESS = "success"
    ERROR = "error"

    @classmethod
    def from_output(cls, output: Output) -> ExecutionStatus:
        match output:
            case StatementOutput():
                return cls.SUCCESS
            case ErrorOutput():
                return cls.ERROR


class StatementOutput(BaseModel):
    """Output from evaluating a Python statement.

    Statements by definition do not return values, just side effects such as
    setting variables or executing functions.

    This is empty for now, but defined to gain type safety (see the `Output`
    tagged union below) and as a holder for any potential future metadata.
    """


class ErrorOutput(BaseModel):
    exception: ExceptionOutput


class ExceptionOutput(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    exc_type: type[Exception]
    message: str
    traceback: TracebackType | None

    @field_serializer("exc_type")
    def serialize_exception_type(self, exc_type: type[Exception]) -> str:
        return exc_type.__name__

    @field_serializer("traceback")
    def serialize_traceback(self, tb: TracebackType | None) -> list[str]:
        if tb is None:
            return []

        tb_lines = traceback.format_tb(tb)
        relevant_tb_lines = [
            line.strip()
            for line in tb_lines
            if "mcp_django/shell" not in line
            and "mcp_django/code" not in line
            and "mcp_django/output" not in line
            and line.strip()
        ]

        return relevant_tb_lines


Output = StatementOutput | ErrorOutput
