from pydantic import BaseModel
from typing import List, Dict, Any, Type

def _recursive_clean(schema: Dict[str, Any], defs: Dict[str, Any]) -> None:
    if not isinstance(schema, dict):
        return

    # 1. Inline $ref if present
    if "$ref" in schema:
        definition_key = schema["$ref"].split("/")[-1]
        if definition_key in defs:
            definition = defs[definition_key]
            schema.update(definition)
            del schema["$ref"]
            
    # 2. Remove title
    schema.pop("title", None)

    # 3. Recurse into properties
    if "properties" in schema:
        for prop_name, prop_val in schema["properties"].items():
            _recursive_clean(prop_val, defs)
    
    # Handle 'items' (arrays)
    if "items" in schema:
        _recursive_clean(schema["items"], defs)

    # Handle complex types
    for key in ["allOf", "anyOf", "oneOf"]:
        if key in schema:
            for item in schema[key]:
                _recursive_clean(item, defs)

def _inject_choice(schema: Dict[str, Any], path_parts: List[str], choices: List[Any]):
    """
    Recursively navigates to a field using dot-notation path and injects enum choices.
    Preserves nullability if the field was originally Optional/Nullable.
    """
    if not path_parts:
        return

    current_field_name = path_parts[0]
    properties = schema.get("properties", {})
    
    if current_field_name not in properties:
        return
        
    field_schema = properties[current_field_name]
    
    if len(path_parts) == 1:
        # Check if the field was nullable (contains 'null' type or in anyOf)
        is_nullable = False
        if "anyOf" in field_schema:
            is_nullable = any(item.get("type") == "null" for item in field_schema["anyOf"])
        elif field_schema.get("type") == "null" or (isinstance(field_schema.get("type"), list) and "null" in field_schema["type"]):
            is_nullable = True

        # Inject enum
        field_schema["enum"] = choices
        
        # Clean up complex types
        field_schema.pop("anyOf", None) 
        field_schema.pop("allOf", None)
        
        # If it was nullable, we must allow null in the enum or the type
        if is_nullable:
            field_schema["type"] = ["string", "null"]
            # To be safe in JSON Schema, we often add None to the enum as well
            if None not in field_schema["enum"]:
                field_schema["enum"].append(None)
        else:
            if not field_schema.get("type"):
                field_schema["type"] = "string"
    else:
        # Recurse deeper
        _inject_choice(field_schema, path_parts[1:], choices)

def get_schema_with_choices(model: Type[BaseModel], dynamic_choices: Dict[str, List[Any]]) -> Dict[str, Any]:
    """
    Generates a clean JSON Schema for a Pydantic model with dynamic enum choices.
    
    This function creates a frontend-friendly JSON Schema by:
    1. Generating the base schema from the Pydantic model
    2. Inlining all $ref definitions to create a flat structure
    3. Removing title fields for cleaner output
    4. Injecting dynamic enum choices into specified fields (supports nested fields via dot notation)
    
    Args:
        model: A Pydantic BaseModel class to generate the schema from
        dynamic_choices: Dictionary mapping field paths to their allowed enum values.
                        Supports dot notation for nested fields.
                        Example: {
                            "region": ["us-east1", "europe-west1"],
                            "address.city": ["Haifa", "Tel Aviv"]
                        }
    
    Returns:
        A clean JSON Schema dictionary with:
        - No $ref or $defs sections (all definitions inlined)
        - No title fields (removed for cleaner output)
        - Injected enum values for specified fields
        - Standard JSON Schema structure compatible with frontend form libraries
    """
    schema = model.model_json_schema()
    defs = schema.get("$defs", {})
    
    # 1. Recursive Clean (Inline Refs + Remove Titles) BEFORE injection
    # This ensures we have a full tree structure to traverse
    _recursive_clean(schema, defs)

    # 2. Inject Dynamic Choices
    for path, choices in dynamic_choices.items():
        _inject_choice(schema, path.split('.'), choices)

    # 3. Cleanup $defs
    if "$defs" in schema:
        del schema["$defs"]

    return schema