"""
Gemini Schema Compatibility Layer

Transforms JSON Schemas generated by Pydantic/FastMCP into
Gemini-compatible format. Applied unconditionally at tool registration.

Why unconditional? Gemini is the "lowest common denominator" — if a schema
works with Gemini, it works with Claude and GPT too. This eliminates
provider-specific bugs without requiring user configuration.

Reference: 
- LOG-001: Root cause ($ref not resolved, degrades to STRING)
- LOG-002: 22 transformation patterns for Gemini compatibility
- LOG-003: Implementation plan (this module)

Maintenance: See docs/GEMINI_SCHEMA_COMPAT.md for update procedures.
"""

import copy
from typing import Any

import jsonref


# Keys that must be removed from schemas (from LOG-002 Section 4 & 5)
_FORBIDDEN_KEYS = frozenset({
    # Not in Gemini Schema spec
    "$id",
    "$schema",
    "additionalProperties",
    "const",  # Will be transformed to enum instead of just removed
    "not",
    "dependentRequired",
    "dependentSchemas",
    "prefixItems",
    "contains",
    "unevaluatedProperties",
    "unevaluatedItems",
    "contentMediaType",
    "contentEncoding",
    "multipleOf",
    "exclusiveMinimum",  # Will be transformed to minimum
    "exclusiveMaximum",  # Will be transformed to maximum
    # Documented but problematic (LOG-002 Section 5)
    "title",
    "default",
    "propertyOrdering",
    "property_ordering",  # snake_case variant
})

# Keys to remove AFTER dereferencing (they contain the definitions)
_DEFS_KEYS = frozenset({"$defs", "definitions"})

# Keys for conditional schemas (not supported by Gemini)
_CONDITIONAL_KEYS = frozenset({"if", "then", "else"})


def make_gemini_compatible(schema: dict) -> dict:
    """
    Transform a JSON Schema to be Gemini-compatible.
    
    This function is called by server.py for every tool schema
    before registration with FastMCP.
    
    Transformations applied (in order):
    1. Dereference all $ref using jsonref
    2. Remove $defs, $id, $schema
    3. Convert anyOf[T, null] → {type: T, nullable: true}
    4. Convert const → enum
    5. Convert exclusiveMinimum/Maximum → minimum/maximum
    6. Remove forbidden keys: title, default, additionalProperties, etc.
    7. Remove conditional schemas: if/then/else
    
    Args:
        schema: Raw JSON Schema from Pydantic model
        
    Returns:
        Gemini-compatible schema (new dict, input not modified)
    """
    result = copy.deepcopy(schema)
    
    # Phase 1: Dereference $ref (CRITICAL - must be first)
    result = _dereference_refs(result)
    
    # Phase 2: Remove $defs/$definitions (after dereferencing)
    result = _remove_defs(result)
    
    # Phase 3: Handle anyOf/oneOf/allOf
    result = _handle_union_types(result)
    
    # Phase 4: Convert const to enum
    result = _convert_const_to_enum(result)
    
    # Phase 5: Handle exclusive bounds
    result = _handle_exclusive_bounds(result)
    
    # Phase 6: Remove forbidden keys
    result = _remove_forbidden_keys(result)
    
    # Phase 7: Remove conditional schemas
    result = _remove_conditional_schemas(result)
    
    return result


def _dereference_refs(schema: dict) -> dict:
    """
    Dereference all $ref using jsonref library.
    
    This is the CRITICAL transformation - without it, nested objects
    degrade to STRING in Gemini (see LOG-001).
    """
    if not _contains_key(schema, "$ref"):
        return schema
    
    # Use jsonref to inline all references
    dereferenced = jsonref.replace_refs(schema, lazy_load=False)
    # Convert back to regular dict (jsonref returns proxy objects)
    return _deep_dict(dereferenced)


def _remove_defs(schema: dict) -> dict:
    """Remove $defs and definitions blocks after dereferencing."""
    result = copy.deepcopy(schema)
    
    for key in _DEFS_KEYS:
        result.pop(key, None)
    
    return result


def _handle_union_types(schema: dict) -> dict:
    """
    Handle anyOf, oneOf, and allOf constructs.
    
    - anyOf with null: Convert to {type: T, nullable: true}
    - oneOf: Convert to anyOf (Gemini only supports anyOf)
    - allOf: Merge schemas
    """
    def process_node(node: Any) -> Any:
        if not isinstance(node, dict):
            return node
        
        result = {}
        
        for key, value in node.items():
            if key == "anyOf" and isinstance(value, list):
                # Check for nullable pattern: anyOf[{type: X}, {type: null}]
                non_null = [v for v in value if not _is_null_type(v)]
                has_null = any(_is_null_type(v) for v in value)
                
                if has_null and len(non_null) == 1:
                    # Convert to nullable
                    inner = process_node(non_null[0])
                    result.update(inner)
                    result["nullable"] = True
                elif len(value) == 1:
                    # Single-item anyOf, just unwrap
                    result.update(process_node(value[0]))
                else:
                    # Keep anyOf but process children
                    result["anyOf"] = [process_node(v) for v in value]
            
            elif key == "oneOf" and isinstance(value, list):
                # Convert oneOf to anyOf (Gemini only supports anyOf)
                result["anyOf"] = [process_node(v) for v in value]
            
            elif key == "allOf" and isinstance(value, list):
                # Merge all schemas in allOf
                merged = {}
                for sub in value:
                    processed = process_node(sub)
                    merged = _merge_schemas(merged, processed)
                result.update(merged)
            
            elif key == "properties" and isinstance(value, dict):
                result["properties"] = {
                    k: process_node(v) for k, v in value.items()
                }
            
            elif key == "items" and isinstance(value, dict):
                result["items"] = process_node(value)
            
            elif isinstance(value, dict):
                result[key] = process_node(value)
            
            elif isinstance(value, list):
                result[key] = [
                    process_node(v) if isinstance(v, dict) else v
                    for v in value
                ]
            
            else:
                result[key] = value
        
        return result
    
    return process_node(schema)


def _convert_const_to_enum(schema: dict) -> dict:
    """Convert const to single-value enum."""
    def process_node(node: Any) -> Any:
        if not isinstance(node, dict):
            return node
        
        result = {}
        for key, value in node.items():
            if key == "const":
                result["enum"] = [value]
            elif isinstance(value, dict):
                result[key] = process_node(value)
            elif isinstance(value, list):
                result[key] = [
                    process_node(v) if isinstance(v, dict) else v
                    for v in value
                ]
            else:
                result[key] = value
        
        return result
    
    return process_node(schema)


def _handle_exclusive_bounds(schema: dict) -> dict:
    """Convert exclusiveMinimum/exclusiveMaximum to minimum/maximum."""
    def process_node(node: Any) -> Any:
        if not isinstance(node, dict):
            return node
        
        result = {}
        for key, value in node.items():
            if key == "exclusiveMinimum":
                result["minimum"] = value
            elif key == "exclusiveMaximum":
                result["maximum"] = value
            elif isinstance(value, dict):
                result[key] = process_node(value)
            elif isinstance(value, list):
                result[key] = [
                    process_node(v) if isinstance(v, dict) else v
                    for v in value
                ]
            else:
                result[key] = value
        
        return result
    
    return process_node(schema)


def _remove_forbidden_keys(schema: dict) -> dict:
    """Remove all forbidden keys recursively."""
    def process_node(node: Any) -> Any:
        if not isinstance(node, dict):
            return node
        
        result = {}
        for key, value in node.items():
            if key in _FORBIDDEN_KEYS:
                continue
            
            if isinstance(value, dict):
                result[key] = process_node(value)
            elif isinstance(value, list):
                result[key] = [
                    process_node(v) if isinstance(v, dict) else v
                    for v in value
                ]
            else:
                result[key] = value
        
        return result
    
    return process_node(schema)


def _remove_conditional_schemas(schema: dict) -> dict:
    """Remove if/then/else constructs."""
    def process_node(node: Any) -> Any:
        if not isinstance(node, dict):
            return node
        
        result = {}
        for key, value in node.items():
            if key in _CONDITIONAL_KEYS:
                continue
            
            if isinstance(value, dict):
                result[key] = process_node(value)
            elif isinstance(value, list):
                result[key] = [
                    process_node(v) if isinstance(v, dict) else v
                    for v in value
                ]
            else:
                result[key] = value
        
        return result
    
    return process_node(schema)


# ============== Helper Functions ==============

def _is_null_type(schema: Any) -> bool:
    """Check if schema represents null type."""
    if not isinstance(schema, dict):
        return False
    return schema.get("type") in ("null", "NULL")


def _contains_key(obj: Any, key: str) -> bool:
    """Recursively check if key exists in nested structure."""
    if isinstance(obj, dict):
        if key in obj:
            return True
        return any(_contains_key(v, key) for v in obj.values())
    elif isinstance(obj, list):
        return any(_contains_key(v, key) for v in obj)
    return False


def _deep_dict(obj: Any) -> Any:
    """Convert jsonref proxy objects back to regular dicts."""
    if isinstance(obj, dict):
        return {k: _deep_dict(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [_deep_dict(v) for v in obj]
    else:
        return obj


def _merge_schemas(base: dict, overlay: dict) -> dict:
    """Merge two schemas for allOf handling."""
    result = copy.deepcopy(base)
    
    for key, value in overlay.items():
        if key == "properties" and "properties" in result:
            result["properties"] = {**result["properties"], **value}
        elif key == "required" and "required" in result:
            result["required"] = list(set(result["required"]) | set(value))
        else:
            result[key] = value
    
    return result