# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utilities for creating optimized Pydantic schemas for LLM usage."""

from typing import Any

from pydantic import BaseModel


class SchemaOptimizer:
    @staticmethod
    def create_optimized_json_schema(model: type[BaseModel]) -> dict[str, Any]:
        """Create the most optimized schema by flattening all $ref/$defs while preserving
        FULL descriptions and ALL action definitions. Also ensures OpenAI strict mode compatibility.

        Args:
                model: The Pydantic model to optimize

        Returns:
                Optimized schema with all $refs resolved and strict mode compatibility
        """
        # Generate original schema
        original_schema = model.model_json_schema()

        # Extract $defs for reference resolution, then flatten everything
        defs_lookup = original_schema.get("$defs", {})

        def optimize_schema(
            obj: Any,
            defs_lookup: dict[str, Any] | None = None,
            *,
            in_properties: bool = False,  # NEW: track context
        ) -> Any:
            """Apply all optimization techniques including flattening all $ref/$defs."""
            if isinstance(obj, dict):
                optimized: dict[str, Any] = {}
                flattened_ref: dict[str, Any] | None = None

                # Skip unnecessary fields AND $defs (we'll inline everything)
                skip_fields = ["additionalProperties", "$defs"]

                for key, value in obj.items():
                    if key in skip_fields:
                        continue

                    # Skip metadata "title" unless we're iterating inside an actual `properties` map
                    if key == "title" and not in_properties:
                        continue

                    # Preserve FULL descriptions without truncation
                    if key in {"description", "type"}:
                        optimized[key] = value

                    # FLATTEN: Resolve $ref by inlining the actual definition
                    elif key == "$ref" and defs_lookup:
                        ref_path = value.split("/")[-1]  # Get the definition name from "#/$defs/SomeName"
                        if ref_path in defs_lookup:
                            # Get the referenced definition and flatten it
                            referenced_def = defs_lookup[ref_path]
                            flattened_ref = optimize_schema(referenced_def, defs_lookup)

                    # Keep all anyOf structures (action unions) and resolve any $refs within
                    elif key == "anyOf" and isinstance(value, list):
                        optimized[key] = [optimize_schema(item, defs_lookup) for item in value]

                    # Recursively optimize nested structures
                    elif key in ["properties", "items"]:
                        optimized[key] = optimize_schema(
                            value,
                            defs_lookup,
                            in_properties=(key == "properties"),
                        )

                    # Keep essential validation fields
                    elif key in [
                        "type",
                        "required",
                        "minimum",
                        "maximum",
                        "minItems",
                        "maxItems",
                        "pattern",
                        "default",
                    ]:
                        optimized[key] = (
                            value if not isinstance(value, dict | list) else optimize_schema(value, defs_lookup)
                        )

                    # Recursively process all other fields
                    else:
                        optimized[key] = (
                            optimize_schema(value, defs_lookup) if isinstance(value, dict | list) else value
                        )

                # If we have a flattened reference, merge it with the optimized properties
                if flattened_ref is not None and isinstance(flattened_ref, dict):
                    # Start with the flattened reference as the base
                    result = flattened_ref.copy()

                    # Merge in any sibling properties that were processed
                    for key, value in optimized.items():
                        # Preserve descriptions from the original object if they exist
                        if (key == "description" and "description" not in result) or key != "description":
                            result[key] = value

                    return result
                # No $ref, just return the optimized object
                # CRITICAL: Add additionalProperties: false to ALL objects for OpenAI strict mode
                if optimized.get("type") == "object":
                    optimized["additionalProperties"] = False

                return optimized

            if isinstance(obj, list):
                return [optimize_schema(item, defs_lookup, in_properties=in_properties) for item in obj]
            return obj

        # Create optimized schema with flattening
        optimized_result = optimize_schema(original_schema, defs_lookup)

        # Ensure we have a dictionary (should always be the case for schema root)
        if not isinstance(optimized_result, dict):
            raise ValueError("Optimized schema result is not a dictionary")

        optimized_schema: dict[str, Any] = optimized_result

        # Additional pass to ensure ALL objects have additionalProperties: false
        def ensure_additional_properties_false(obj: Any) -> None:
            """Ensure all objects have additionalProperties: false."""
            if isinstance(obj, dict):
                # If it's an object type, ensure additionalProperties is false
                if obj.get("type") == "object":
                    obj["additionalProperties"] = False

                # Recursively apply to all values
                for value in obj.values():
                    if isinstance(value, dict | list):
                        ensure_additional_properties_false(value)
            elif isinstance(obj, list):
                for item in obj:
                    if isinstance(item, dict | list):
                        ensure_additional_properties_false(item)

        ensure_additional_properties_false(optimized_schema)
        SchemaOptimizer._make_strict_compatible(optimized_schema)

        return optimized_schema

    @staticmethod
    def _make_strict_compatible(schema: dict[str, Any] | list[Any]) -> None:
        """Ensure all properties are required for OpenAI strict mode."""
        if isinstance(schema, dict):
            # First recursively apply to nested objects
            for key, value in schema.items():
                if isinstance(value, dict | list) and key != "required":
                    SchemaOptimizer._make_strict_compatible(value)

            # Then update required for this level
            if "properties" in schema and "type" in schema and schema["type"] == "object":
                # Add all properties to required array
                all_props = list(schema["properties"].keys())
                schema["required"] = all_props  # Set all properties as required

        elif isinstance(schema, list):
            for item in schema:
                SchemaOptimizer._make_strict_compatible(item)
