"""
OpenCode Style agent 的 schema 合并与响应格式工具。

本模块从 `factory.py` 中拆分出以下与图结构无关的纯工具逻辑：
- 基于多份 TypedDict schema 合并出 state / input / output schema
- 处理 `OmitFromSchema` / `Annotated` 等元数据，控制字段是否出现在输入或输出中
- 探测当前模型是否支持 provider 侧结构化输出（provider strategy）
- 处理 structured output 错误并决定是否重试，以及给模型的错误提示内容

这些工具保持与 LangChain 官方语义一致，可在不同 agent 工厂中复用。
"""

from __future__ import annotations

from typing import Any, get_args, get_origin, get_type_hints

from langchain.agents.middleware.types import OmitFromSchema
from langchain.agents.structured_output import (
    AutoStrategy,
    MultipleStructuredOutputsError,
    OutputToolBinding,
    ProviderStrategy,
    ProviderStrategyBinding,
    ResponseFormat,
    StructuredOutputError,
    StructuredOutputValidationError,
    ToolStrategy,
)
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.tools import BaseTool
from typing_extensions import Annotated, NotRequired, Required, TypedDict


STRUCTURED_OUTPUT_ERROR_TEMPLATE = "Error: {error}\n Please fix your mistakes."

FALLBACK_MODELS_WITH_STRUCTURED_OUTPUT = [
    # if model profile data are not available, these models are assumed to support
    # structured output
    "grok",
    "gpt-5",
    "gpt-4.1",
    "gpt-4o",
    "gpt-oss",
    "o3-pro",
    "o3-mini",
]


def _resolve_schema(schemas: set[type], schema_name: str, omit_flag: str | None = None) -> type:
    """Resolve schema by merging schemas and optionally respecting `OmitFromSchema` annotations."""
    all_annotations: dict[str, Any] = {}

    for schema in schemas:
        hints = get_type_hints(schema, include_extras=True)

        for field_name, field_type in hints.items():
            should_omit = False

            if omit_flag:
                metadata = _extract_metadata(field_type)
                for meta in metadata:
                    if isinstance(meta, OmitFromSchema) and getattr(meta, omit_flag) is True:
                        should_omit = True
                        break

            if not should_omit:
                all_annotations[field_name] = field_type

    return TypedDict(schema_name, all_annotations)  # type: ignore[operator]


def _extract_metadata(type_: type) -> list[Any]:
    """Extract metadata from a field type, handling `Required`/`NotRequired` and `Annotated` wrappers."""
    if get_origin(type_) in {Required, NotRequired}:
        inner_type = get_args(type_)[0]
        if get_origin(inner_type) is Annotated:
            return list(get_args(inner_type)[1:])

    if get_origin(type_) is Annotated:
        return list(get_args(type_)[1:])

    return []


def _supports_provider_strategy(
    model: str | BaseChatModel, tools: list[BaseTool | dict[str, Any]] | None = None
) -> bool:
    """Check if a model supports provider-specific structured output."""
    model_name: str | None = None
    if isinstance(model, str):
        model_name = model
    elif isinstance(model, BaseChatModel):
        model_name = (
            getattr(model, "model_name", None)
            or getattr(model, "model", None)
            or getattr(model, "model_id", "")
        )
        model_profile = model.profile
        if (
            model_profile is not None
            and model_profile.get("structured_output")
            and not (
                tools
                and isinstance(model_name, str)
                and "gemini" in model_name.lower()
                and "gemini-3" not in model_name.lower()
            )
        ):
            return True

    return (
        any(part in model_name.lower() for part in FALLBACK_MODELS_WITH_STRUCTURED_OUTPUT)
        if model_name
        else False
    )


def _handle_structured_output_error(
    exception: Exception,
    response_format: ResponseFormat[Any],
) -> tuple[bool, str]:
    """Handle structured output error. Returns `(should_retry, retry_tool_message)`."""
    if not isinstance(response_format, ToolStrategy):
        return False, ""

    handle_errors = response_format.handle_errors

    if handle_errors is False:
        return False, ""
    if handle_errors is True:
        return True, STRUCTURED_OUTPUT_ERROR_TEMPLATE.format(error=str(exception))
    if isinstance(handle_errors, str):
        return True, handle_errors
    if isinstance(handle_errors, type):
        if issubclass(handle_errors, Exception) and isinstance(exception, handle_errors):
            return True, STRUCTURED_OUTPUT_ERROR_TEMPLATE.format(error=str(exception))
        return False, ""
    if isinstance(handle_errors, tuple):
        if any(isinstance(exception, exc_type) for exc_type in handle_errors):
            return True, STRUCTURED_OUTPUT_ERROR_TEMPLATE.format(error=str(exception))
        return False, ""
    return True, handle_errors(exception)


__all__ = [
    "STRUCTURED_OUTPUT_ERROR_TEMPLATE",
    "FALLBACK_MODELS_WITH_STRUCTURED_OUTPUT",
    "_resolve_schema",
    "_extract_metadata",
    "_supports_provider_strategy",
    "_handle_structured_output_error",
]

