from __future__ import annotations

import logging
from collections.abc import Generator, Sequence
from contextlib import contextmanager
from typing import TYPE_CHECKING, TypeAlias, cast, get_type_hints

from ...content import Text
from ...messages import AssistantMessage, Message, SystemMessage, UserMessage
from ...models.params import (
    Params,  # Import directly from params.py to avoid circular dependency
)
from ...tools import AnyToolSchema, ProviderTool, ToolSchema
from ..provider_id import ProviderId

if TYPE_CHECKING:
    from ...models import ThinkingConfig
    from ..model_id import ModelId

logger = logging.getLogger(__name__)

SystemMessageContent: TypeAlias = str | None


def get_include_thoughts(params: Params) -> bool:
    """Extract include_thoughts from params thinking config."""
    thinking_config = params.get("thinking")
    return (thinking_config or {}).get("include_thoughts", False)


def has_strict_tools(tools: Sequence[AnyToolSchema | ProviderTool]) -> bool:
    """Check if any tools have strict=True explicitly set.

    Args:
        tools: The tools to check

    Returns:
        True if any tool has strict=True, False otherwise
    """
    return any(isinstance(tool, ToolSchema) and tool.strict is True for tool in tools)


def ensure_additional_properties_false(obj: object) -> None:
    """Recursively adds additionalProperties = False to a schema, required for strict mode."""
    if isinstance(obj, dict):
        obj = cast(dict[str, object], obj)
        if obj.get("type") == "object" and "additionalProperties" not in obj:
            obj["additionalProperties"] = False
        for value in obj.values():
            ensure_additional_properties_false(value)
    elif isinstance(obj, list):
        obj = cast(list[object], obj)
        for item in obj:
            ensure_additional_properties_false(item)


def ensure_all_properties_required(obj: object) -> None:
    """Recursively ensures all properties are in required array, needed for OpenAI strict mode.

    OpenAI's strict mode requires that all properties in an object schema are listed
    in the 'required' array, even if they have default values.
    """
    if isinstance(obj, dict):
        obj = cast(dict[str, object], obj)
        if obj.get("type") == "object" and "properties" in obj:
            properties = obj.get("properties")
            if isinstance(properties, dict):
                property_keys = cast(dict[str, object], properties)
                obj["required"] = list(property_keys.keys())
        for value in obj.values():
            ensure_all_properties_required(value)
    elif isinstance(obj, list):
        obj = cast(list[object], obj)
        for item in obj:
            ensure_all_properties_required(item)


def add_system_instructions(
    messages: Sequence[Message], additional_system_instructions: str
) -> Sequence[Message]:
    """Add system instructions to a sequence of messages.

    If the first message is a system message, appends the additional instructions
    to it with a newline separator. If the instructions already exist at the end
    of the system message, returns the original messages unchanged. Otherwise,
    creates a new system message at the beginning of the sequence.

    Args:
        messages: The sequence of messages to modify.
        additional_system_instructions: The system instructions to add.

    Returns:
        A new sequence of messages with the system instructions added.
    """
    if messages and messages[0].role == "system":
        if messages[0].content.text.endswith(additional_system_instructions):
            return messages
        modified = Text(
            text=messages[0].content.text + "\n" + additional_system_instructions
        )
        return [SystemMessage(role="system", content=modified), *messages[1:]]
    else:
        return [
            SystemMessage(
                role="system", content=Text(text=additional_system_instructions)
            ),
            *messages,
        ]


def extract_system_message(
    messages: Sequence[Message],
) -> tuple[SystemMessageContent, Sequence[UserMessage | AssistantMessage]]:
    """Extract the system message(s) from a list of Messages.

    This takes a list of messages, and returns the list of messages with
    all system messages removed, as well as the textual contents of the first message,
    if that message was a system message. If there are any system messages that are
    not the first message, they will be dropped, and a warning will be emitted.

    This is intended for use in clients where the system message is not included in the
    input messages, but passed as an additional argument or metadata.
    """
    system_message_content: SystemMessageContent = None
    remaining_messages: list[UserMessage | AssistantMessage] = []

    for i, message in enumerate(messages):
        if message.role == "system":
            if i == 0:
                system_message_content = message.content.text
            else:
                logging.warning(
                    "Skipping system message at index %d because it is not the first message",
                    i,
                )
        else:
            remaining_messages.append(message)

    return system_message_content, remaining_messages


class SafeParamsAccessor:
    """A wrapper around Params that tracks which parameters have been accessed."""

    def __init__(self, params: Params) -> None:
        self._params = params
        self._unaccessed = set(get_type_hints(Params).keys())

    @property
    def temperature(self) -> float | None:
        """Access the temperature parameter."""
        self._unaccessed.discard("temperature")
        return self._params.get("temperature")

    @property
    def max_tokens(self) -> int | None:
        """Access the max_tokens parameter."""
        self._unaccessed.discard("max_tokens")
        return self._params.get("max_tokens")

    @property
    def top_p(self) -> float | None:
        """Access the top_p parameter."""
        self._unaccessed.discard("top_p")
        return self._params.get("top_p")

    @property
    def top_k(self) -> int | None:
        """Access the top_k parameter."""
        self._unaccessed.discard("top_k")
        return self._params.get("top_k")

    @property
    def seed(self) -> int | None:
        """Access the seed parameter."""
        self._unaccessed.discard("seed")
        return self._params.get("seed")

    @property
    def stop_sequences(self) -> list[str] | None:
        """Access the stop_sequences parameter."""
        self._unaccessed.discard("stop_sequences")
        return self._params.get("stop_sequences")

    @property
    def thinking(self) -> ThinkingConfig | None:
        """Access the thinking parameter."""
        self._unaccessed.discard("thinking")
        return self._params.get("thinking")

    def emit_warning_for_unused_param(
        self,
        param_name: str,
        param_value: object,
        provider_id: ProviderId,
        model_id: ModelId | None = None,
    ) -> None:
        unsupported_by = f"provider: {provider_id}"
        if model_id:
            unsupported_by += f" with model_id: {model_id}"
        logger.warning(
            f"Skipping unsupported parameter: {param_name}={param_value} ({unsupported_by})"
        )

    def check_access_integrity(self, unsupported_params: list[str]) -> None:
        """Verify that all used parameters have been accessed, and none of the unsupported have been."""
        assert self._unaccessed == set(unsupported_params), (
            "Mismatch between unsupported and unaccessed params"
        )


@contextmanager
def ensure_all_params_accessed(
    *,
    params: Params,
    provider_id: ProviderId,
    unsupported_params: list[str] | None = None,
) -> Generator[SafeParamsAccessor, None, None]:
    """Context manager that ensures all parameters are accessed.

    Yields a wrapper around params that tracks which parameters have been accessed.
    On context exit, raises a `RuntimeError` if any parameters were not accessed.

    Args:
        params: The parameters to wrap
        provider: The provider that is accessing these params (required for logging)
        unsupported_params: A list of params keys it does not support, for auto-warning
            of unsupported params as boilerplate reduction. Or None, to disable this
            optional feature.

    Yields:
        A SafeParamsAccessor instance if params is not None, else None

    Raises:
        RuntimeError: If any parameters were not accessed before context exit
    """

    accessor = SafeParamsAccessor(params)
    unsupported_params = unsupported_params or []
    for unsupported in unsupported_params:
        if (val := params.get(unsupported)) is not None:
            accessor.emit_warning_for_unused_param(
                unsupported, val, provider_id=provider_id
            )
    try:
        yield accessor
    finally:
        accessor.check_access_integrity(unsupported_params=unsupported_params)
