from __future__ import annotations

import inspect
import json
import os
import warnings
from typing import Any, Generic, Literal, Optional, Sequence, TypeVar, cast

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

from ._content import (
    Content,
    ContentText,
    ContentToolRequest,
    ContentToolResult,
    ContentUnion,
    create_content,
)
from ._content_expand import expand_tool_result, merge_tool_results, tool_results_first

__all__ = ("Turn", "UserTurn", "SystemTurn", "AssistantTurn")


CompletionT = TypeVar("CompletionT")
Role = Literal["user", "assistant", "system"]

# The reasons chatlas normalizes provider-specific finish/stop reasons to.
# Not exhaustive: providers may surface a reason chatlas doesn't yet
# recognize, so usages pair this with `| str` (see AssistantTurn.finish_reason).
FinishReason = Literal[
    "success",
    "tool_use",
    "max_tokens",
    "content_filter",
    "context_window",
    "stop_sequence",
]


class Turn(BaseModel):
    """
    Base turn class

    Every conversation with a chatbot consists of pairs of user and assistant
    turns, corresponding to an HTTP request and response. These turns are
    represented by `Turn` objects (or their subclasses `UserTurn`, `SystemTurn`,
    `AssistantTurn`), which contain a list of [](`~chatlas.types.Content`)s
    representing the individual messages within the turn. These might be text,
    images, tool requests (assistant only), or tool responses (user only).

    Note that a call to `.chat()` and related functions may result in multiple
    user-assistant turn cycles. For example, if you have registered tools, chatlas
    will automatically handle the tool calling loop, which may result in any
    number of additional cycles.

    Examples
    --------

    ```python
    from chatlas import UserTurn, AssistantTurn, ChatOpenAI, ChatAnthropic

    chat = ChatOpenAI()
    str(chat.chat("What is the capital of France?"))
    turns = chat.get_turns()
    assert len(turns) == 2
    assert isinstance(turns[0], UserTurn)
    assert turns[0].role == "user"
    assert isinstance(turns[1], AssistantTurn)
    assert turns[1].role == "assistant"

    # Load context into a new chat instance
    chat2 = ChatAnthropic()
    chat2.set_turns(turns)
    turns2 = chat2.get_turns()
    assert turns == turns2
    ```

    Parameters
    ----------
    contents
        A list of [](`~chatlas.types.Content`) objects.
    """

    contents: list[ContentUnion] = Field(default_factory=list)

    # Discriminator field for Pydantic to determine which subclass to instantiate
    role: Role
    model_config = ConfigDict(arbitrary_types_allowed=True)

    def __init__(
        self,
        contents: str | Sequence[Content | str],
        **kwargs,
    ):
        if isinstance(contents, str):
            contents = [ContentText(text=contents)]

        contents2: list[Content] = []
        for x in contents:
            if isinstance(x, Content):
                contents2.append(x)
            elif isinstance(x, str):
                contents2.append(ContentText(text=x))
            elif isinstance(x, dict):
                contents2.append(create_content(x))
            else:
                raise ValueError("All contents must be Content objects or str.")

        super().__init__(
            contents=contents2,
            **kwargs,
        )

    @property
    def text(self) -> str:
        return "".join(x.text for x in self.contents if isinstance(x, ContentText))

    def __str__(self) -> str:
        return "\n".join(str(c) for c in self.contents)

    def __repr__(self):
        return self.__str__()

    def _repr_markdown_(self):
        return self.__str__()

    @classmethod
    def model_validate(
        cls,
        obj: Any,
        *,
        strict: Optional[bool] = None,
        extra: Any = None,
        from_attributes: Optional[bool] = None,
        context: Any = None,
        by_alias: Optional[bool] = None,
        by_name: Optional[bool] = None,
    ) -> "Turn":
        if cls is not Turn:
            # If called on a subclass, use the standard validation
            return super().model_validate(
                obj,
                strict=strict,
                extra=extra,
                from_attributes=from_attributes,
                context=context,
                by_alias=by_alias,
                by_name=by_name,
            )

        # Determine the correct subclass based on the role field
        if isinstance(obj, dict) and "role" in obj:
            target_cls = cls._get_turn_class_for_role(obj["role"])
            return target_cls.model_validate(
                obj,
                strict=strict,
                extra=extra,
                from_attributes=from_attributes,
                context=context,
                by_alias=by_alias,
                by_name=by_name,
            )

        # Fallback to default behavior
        return super().model_validate(
            obj,
            strict=strict,
            extra=extra,
            from_attributes=from_attributes,
            context=context,
            by_alias=by_alias,
            by_name=by_name,
        )

    @classmethod
    def model_validate_json(
        cls,
        json_data: str | bytes | bytearray,
        *,
        strict: Optional[bool] = None,
        extra: Any = None,
        context: Any = None,
        by_alias: Optional[bool] = None,
        by_name: Optional[bool] = None,
    ) -> "Turn":
        if cls is not Turn:
            # If called on a subclass, use the standard validation
            return super().model_validate_json(
                json_data,
                strict=strict,
                extra=extra,
                context=context,
                by_alias=by_alias,
                by_name=by_name,
            )

        # Parse JSON to determine the role
        obj = json.loads(json_data)
        if isinstance(obj, dict) and "role" in obj:
            target_cls = cls._get_turn_class_for_role(obj["role"])
            return target_cls.model_validate_json(
                json_data,
                strict=strict,
                extra=extra,
                context=context,
                by_alias=by_alias,
                by_name=by_name,
            )
        else:
            return super().model_validate_json(
                json_data,
                strict=strict,
                extra=extra,
                context=context,
                by_alias=by_alias,
                by_name=by_name,
            )

    @classmethod
    def _get_turn_class_for_role(cls, role: str) -> type["Turn"]:
        if role == "user":
            return UserTurn
        elif role == "system":
            return SystemTurn
        elif role == "assistant":
            return AssistantTurn
        else:
            raise ValueError(f"Unknown role: {role}")

    def to_inspect_messages(self, model: Optional[str] = None):
        """
        Transform this turn into a list of Inspect AI `ChatMessage` objects.

        Most users will not need to call this method directly. See the
        `.export_eval()` method on `Chat` for a higher level interface to
        exporting chat history for evaluation purposes.
        """

        from ._inspect import try_import_inspect, turn_as_inspect_messages

        try_import_inspect()
        return turn_as_inspect_messages(self, model=model)


class UserTurn(Turn):
    """
    User turn - represents user input

    Parameters
    ----------
    contents
        A list of [](`~chatlas.types.Content`) objects, or strings.

    See Also
    --------
    - :class:`~chatlas.Turn`: The base class for all turn types.
    """

    role: Literal["user"] = Field(default="user", frozen=True)  # pyright: ignore[reportIncompatibleVariableOverride]

    # Make contents a positional argument for convenience
    def __init__(
        self,
        contents: str | Sequence[Content | str],
        **kwargs,
    ):
        super().__init__(contents, **kwargs)

    @model_validator(mode="after")
    def expand_tool_contents(self):
        contents: list[ContentUnion] = []
        for x in merge_tool_results(self.contents):
            if isinstance(x, ContentToolResult):
                contents.extend(expand_tool_result(x))
            else:
                contents.append(x)

        self.contents = tool_results_first(contents)
        return self


class SystemTurn(Turn):
    """
    System turn - represents system prompt

    Parameters
    ----------
    contents
        A list of [](`~chatlas.types.Content`) objects, or strings.

    See Also
    --------
    - :class:`~chatlas.Turn`: The base class for all turn types.
    """

    role: Literal["system"] = Field(default="system", frozen=True)  # pyright: ignore[reportIncompatibleVariableOverride]

    # Make contents a positional argument for convenience
    def __init__(
        self,
        contents: str | Sequence[Content | str],
        **kwargs,
    ):
        super().__init__(contents, **kwargs)


class AssistantTurn(Turn, Generic[CompletionT]):
    """
    Assistant turn - represents model response with additional metadata

    Parameters
    ----------
    contents
        A list of [](`~chatlas.types.Content`) objects.
    tokens
        A numeric vector of length 3 representing the number of input, output, and cached
        tokens (respectively) used in this turn.
    finish_reason
        The (normalized) reason why the turn ended, e.g. `"success"`,
        `"tool_use"`, `"max_tokens"`, `"content_filter"`, `"context_window"`,
        or `"stop_sequence"`. Providers may also surface a reason chatlas
        doesn't yet recognize, in which case the raw string is passed through
        unchanged.
    completion
        The completion object returned by the provider. This is useful if there's
        information returned by the provider that chatlas doesn't otherwise expose.
    cost
        The cost of this turn in USD. This is computed when the turn is created
        based on the token usage and pricing information (including service tier).
    partial_reason
        If set, indicates this turn is incomplete (e.g., the stream was interrupted
        or cancelled). The value describes the reason for the partial state.

    See Also
    --------
    - :class:`~chatlas.Turn`: The base class for all turn types.
    """

    role: Literal["assistant"] = Field(  # pyright: ignore[reportIncompatibleVariableOverride]
        default="assistant",
        frozen=True,
    )
    tokens: Optional[tuple[int, int, int]] = None
    finish_reason: Optional[FinishReason | str] = None
    completion: Optional[CompletionT] = Field(default=None, exclude=True)
    cost: Optional[float] = None
    partial_reason: Optional[str] = None

    @property
    def is_partial(self) -> bool:
        """Whether this turn is a partial (interrupted/cancelled) turn."""
        return self.partial_reason is not None

    @field_validator("tokens", mode="before")
    @classmethod
    def validate_tokens(cls, v):
        """Convert list to tuple for JSON deserialization compatibility."""
        if isinstance(v, list):
            return tuple(v)
        return v

    def __init__(
        self,
        contents: str | Sequence[Content | str],
        *,
        tokens: Optional[tuple[int, int, int] | list[int]] = None,
        finish_reason: Optional[FinishReason | str] = None,
        completion: Optional[CompletionT] = None,
        cost: Optional[float] = None,
        partial_reason: Optional[str] = None,
        **kwargs,
    ):
        if isinstance(tokens, list):
            tokens = cast(tuple[int, int, int], tuple(tokens))

        # Pass assistant-specific fields to parent constructor
        if tokens is not None:
            kwargs["tokens"] = tokens
        if finish_reason is not None:
            kwargs["finish_reason"] = finish_reason
        if completion is not None:
            kwargs["completion"] = completion
        if cost is not None:
            kwargs["cost"] = cost
        if partial_reason is not None:
            kwargs["partial_reason"] = partial_reason

        super().__init__(contents, **kwargs)


def user_turn(
    *args: Content | str,
    prior_turns: Optional[list[Turn]] = None,
) -> UserTurn:
    if len(args) == 0:
        raise ValueError("Must supply at least one input.")

    # If the prior turns contain dangling tool requests
    # (possibly due to an interrupted chat), then complete them.
    results = complete_dangling_tool_requests(prior_turns or [])

    if not results:
        return UserTurn(args)

    # Include the tool results as additional contents in the user turn
    # Tool results must come first, before the new user input
    contents = results + list(args)
    return UserTurn(contents)


def complete_dangling_tool_requests(turns: Sequence[Turn]) -> list[ContentToolResult]:
    if len(turns) == 0:
        return []

    last_turn = turns[-1]
    if not isinstance(last_turn, AssistantTurn):
        return []

    tool_requests = [
        content
        for content in last_turn.contents
        if isinstance(content, ContentToolRequest)
    ]
    if not tool_requests:
        return []

    return [
        ContentToolResult(
            value=None,
            error=ToolNotInvokedError("Chat ended before the tool could be invoked."),
            request=req,
        )
        for req in tool_requests
    ]


def check_finish_reason(
    finish_reason: FinishReason | str | None,
    signal: Literal["error", "warn"],
) -> None:
    """
    Signal that a response finished for a reason that leaves it incomplete.

    Use `signal="error"` where nothing useful can be done with a partial response
    and `signal="warn"` where it's still worth handing back. Structured data
    extraction is the former: truncated JSON fails to parse, and the resulting
    decode error says nothing about the cause. So providers call this with
    `"error"` *before* parsing a data model out of the response (gh-315), leaving
    `Chat` to `"warn"` about whatever partial responses do get through.
    """
    msg = INCOMPLETE_FINISH_REASONS.get(finish_reason) if finish_reason else None
    if msg is None:
        return

    if signal == "error":
        raise ValueError(msg)
    warnings.warn(msg, UserWarning, stacklevel=caller_stacklevel())


def caller_stacklevel() -> int:
    """
    The `warnings.warn()` stacklevel of the nearest frame outside chatlas.

    A fixed stacklevel can't work here: the distance from a warning to the user's
    own code depends on how they got there (`.chat()` sits two frames deeper than
    `.stream()`, and both move again under `parallel_chat()`), so a constant
    tuned for one entry point misattributes every other one. Python 3.12's
    `warnings.warn(skip_file_prefixes=)` does exactly this, but chatlas supports
    3.10+.
    """
    package_dir = os.path.dirname(os.path.abspath(__file__)) + os.sep

    frame = inspect.currentframe()
    frame = frame.f_back if frame else None  # start at our caller
    level = 1
    while frame is not None:
        if not os.path.abspath(frame.f_code.co_filename).startswith(package_dir):
            return level
        frame = frame.f_back
        level += 1

    # Nothing outside chatlas on the stack (e.g. a worker thread); blame the
    # warn() site rather than overshooting the top of the stack.
    return 1


# Finish reasons that mean the response isn't the one the model set out to give.
# Keyed by the normalized reason (see `FinishReason`).
INCOMPLETE_FINISH_REASONS: dict[str, str] = {
    "max_tokens": (
        "Response was truncated because it hit the `max_tokens` limit. "
        "Increase `max_tokens` to allow the model to generate the full response."
    ),
    "context_window": (
        "Response was truncated because it exceeded the model's context window."
    ),
    "content_filter": (
        "Response was filtered by the provider's content moderation policy."
    ),
}


class ToolNotInvokedError(Exception):
    """Raised when a tool was requested but not invoked before chat ended."""

    pass
