from typing import Annotated, Any, Dict, List, Optional, Set, Type, TypeVar

from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter, create_model

from . import SchemaError
from .builders import ConstraintBuilder
from .handlers import CombinerHandler
from .interfaces import IModelBuilder
from .resolvers import ReferenceResolver, TypeResolver

T = TypeVar("T", bound=BaseModel)


class PydanticModelBuilder(IModelBuilder[T]):
    """Creates Pydantic models from JSON Schema definitions"""

    # Standard JSON Schema properties for fields
    STANDARD_FIELD_PROPERTIES = {
        "type",
        "format",
        "description",
        "default",
        "title",
        "examples",
        "const",
        "enum",
        "multipleOf",
        "maximum",
        "exclusiveMaximum",
        "minimum",
        "exclusiveMinimum",
        "maxLength",
        "minLength",
        "pattern",
        "items",
        "additionalItems",
        "maxItems",
        "minItems",
        "uniqueItems",
        "properties",
        "additionalProperties",
        "required",
        "patternProperties",
        "dependencies",
        "propertyNames",
        "if",
        "then",
        "else",
        "allOf",
        "anyOf",
        "oneOf",
        "not",
        "$ref",
        "$defs",
        "definitions",
    }

    # Standard JSON Schema properties for models
    STANDARD_MODEL_PROPERTIES = {
        "type",
        "title",
        "description",
        "properties",
        "required",
        "additionalProperties",
        "patternProperties",
        "dependencies",
        "propertyNames",
        "if",
        "then",
        "else",
        "allOf",
        "anyOf",
        "oneOf",
        "not",
        "$ref",
        "$defs",
        "definitions",
        "$schema",
        "$id",
        "$comment",
        "items",
        "minItems",
        "maxItems",
        "uniqueItems",
    }

    def __init__(
        self,
        base_model_type: Type[T] = BaseModel,
        predefined_models: Optional[Dict[str, Type[BaseModel]]] = None,
        predefined_refs: Optional[Dict[str, Any]] = None,
    ):
        # Instantiate resolvers and builders directly
        self.type_resolver = TypeResolver()
        self.constraint_builder = ConstraintBuilder()
        self.reference_resolver = ReferenceResolver()
        # Pass resolvers and method references as callbacks to CombinerHandler
        self.combiner_handler = CombinerHandler(
            type_resolver=self.type_resolver,
            constraint_builder=self.constraint_builder,
            reference_resolver=self.reference_resolver,
            recursive_field_builder=self._get_field_type,
            field_info_builder=self._build_field_info,
            name_sanitizer=self._sanitize_field_name,
        )
        self.base_model_type = base_model_type
        validated_predefined_models = self._validate_predefined_models(
            predefined_models,
            base_model_type=self.base_model_type,
        )
        validated_predefined_refs = self._validate_predefined_refs(predefined_refs)
        overlapping_refs = set(validated_predefined_models).intersection(validated_predefined_refs)
        if overlapping_refs:
            overlap_list = ", ".join(sorted(overlapping_refs))
            raise ValueError(
                "Duplicate predefined refs found in both predefined_models and predefined_refs: "
                f"{overlap_list}. Use only one mapping per ref."
            )
        # Track models being built to handle recursive references
        self._model_cache: Dict[str, Type[BaseModel]] = dict(validated_predefined_models)
        self._ref_type_cache: Dict[str, Any] = dict(validated_predefined_refs)
        self._building_models: Set[str] = set()
        self._models_to_rebuild: Set[Type[BaseModel]] = set()

    @staticmethod
    def _validate_ref_key(ref: Any, mapping_name: str) -> str:
        """Validate local JSON Pointer reference keys used in predefined mappings."""
        if not isinstance(ref, str) or not ref.startswith("#/"):
            raise ValueError(
                f"Invalid {mapping_name} ref '{ref}'. Keys must be local JSON Pointer refs like "
                "'#/definitions/Model'"
            )

        path = ref[2:]
        if not path or any(segment == "" for segment in path.split("/")):
            raise ValueError(
                f"Invalid {mapping_name} ref '{ref}'. Keys must be local JSON Pointer refs "
                "without empty path segments, for example '#/definitions/Model'"
            )
        return ref

    def _validate_predefined_models(
        self,
        predefined_models: Optional[Dict[str, Type[BaseModel]]],
        base_model_type: Type[T],
    ) -> Dict[str, Type[BaseModel]]:
        """Validate and normalize predefined model mappings for $ref resolution."""
        if predefined_models is None:
            return {}
        if not isinstance(predefined_models, dict):
            raise ValueError(
                "predefined_models must be a dict mapping local $ref strings to BaseModel subclasses"
            )

        validated: Dict[str, Type[BaseModel]] = {}
        for ref, model in predefined_models.items():
            ref = self._validate_ref_key(ref, "predefined model")

            if not isinstance(model, type) or not issubclass(model, BaseModel):
                raise ValueError(
                    f"Invalid predefined model for ref '{ref}'. Values must be subclasses of pydantic.BaseModel"
                )
            if not issubclass(model, base_model_type):
                raise ValueError(
                    f"Invalid predefined model for ref '{ref}'. Values must be subclasses of the configured "
                    f"base_model_type ({base_model_type.__name__})"
                )
            validated[ref] = model
        return validated

    def _validate_predefined_refs(
        self,
        predefined_refs: Optional[Dict[str, Any]],
    ) -> Dict[str, Any]:
        """Validate and normalize predefined non-model reference mappings."""
        if predefined_refs is None:
            return {}
        if not isinstance(predefined_refs, dict):
            raise ValueError(
                "predefined_refs must be a dict mapping local $ref strings to valid type annotations"
            )

        validated: Dict[str, Any] = {}
        for ref, annotation in predefined_refs.items():
            ref = self._validate_ref_key(ref, "predefined ref")
            try:
                TypeAdapter(annotation)
            except Exception as exc:
                raise ValueError(
                    f"Invalid predefined ref type for '{ref}'. Value must be a valid "
                    "Pydantic-compatible type annotation."
                ) from exc
            validated[ref] = annotation
        return validated

    @staticmethod
    def _get_ref_title(ref: str) -> str:
        ref_parts = ref.split("/")
        return ref_parts[-1] if ref_parts else "DynamicModel"

    def _create_predefined_type_root_model(self, ref: str) -> Type[T]:
        """Wrap a predefined non-model ref type into a RootModel for top-level $ref schemas."""
        value_type = self._ref_type_cache[ref]
        title = self._get_ref_title(ref)
        return type(title, (RootModel[value_type],), {})

    def create_pydantic_model(
        self,
        schema: Dict[str, Any],
        root_schema: Optional[Dict[str, Any]] = None,
        allow_undefined_array_items: bool = False,
        allow_undefined_type: bool = False,
        populate_by_name: bool = False,
        _schema_ref: Optional[str] = None,
    ) -> Type[T]:
        """
        Creates a Pydantic model from a JSON Schema definition.

        Args:
            schema: The JSON Schema to convert
            root_schema: The root schema containing definitions
            allow_undefined_array_items: If True, allows arrays without items schema
            allow_undefined_type: If True, allows schemas without an explicit type
            populate_by_name: If True, allows population of model fields by name and alias
            _schema_ref: Internal parameter to track the original $ref of the schema

        Returns:
            A Pydantic model class
        """
        if root_schema is None:
            root_schema = schema

        # Store the original ref if present for tracking
        # Use the passed ref or extract from schema
        original_ref = _schema_ref or schema.get("$ref")

        # Handle references
        if "$ref" in schema:
            # Check if we've already built this model
            if original_ref in self._model_cache:
                return self._model_cache[original_ref]
            if original_ref in self._ref_type_cache:
                predefined_model = self._create_predefined_type_root_model(original_ref)
                self._model_cache[original_ref] = predefined_model
                return predefined_model

            # Mark this ref as being built
            if original_ref:
                self._building_models.add(original_ref)

            schema = self.reference_resolver.resolve_ref(
                schema["$ref"],
                schema,
                root_schema,
            )

        # Handle combiners
        if "allOf" in schema:
            return self.combiner_handler.handle_all_of(
                schema["allOf"],
                root_schema,
                allow_undefined_array_items,
                allow_undefined_type,
                populate_by_name,
            )
        if "anyOf" in schema:
            return self.combiner_handler.handle_any_of(
                schema["anyOf"],
                root_schema,
                allow_undefined_array_items,
                allow_undefined_type,
                populate_by_name,
            )
        if "oneOf" in schema:
            return self.combiner_handler.handle_one_of(
                schema["oneOf"],
                root_schema,
                allow_undefined_array_items,
                allow_undefined_type,
                populate_by_name
            )

        # Handle top-level arrays
        if schema.get("type") == "array":
            return self._create_array_root_model(
                schema, root_schema, allow_undefined_array_items, original_ref
            )

        # Handle top-level scalars
        # We check for explicit scalar types or enum/const
        schema_type = schema.get("type")
        is_scalar = (
            (
                isinstance(schema_type, str)
                and schema_type in ("string", "integer", "number", "boolean", "null")
            )
            or (
                isinstance(schema_type, list)
                and "object" not in schema_type
                and "array" not in schema_type
            )
            or "enum" in schema
            or "const" in schema
        )

        if is_scalar:
            return self._create_scalar_root_model(
                schema,
                root_schema,
                allow_undefined_array_items,
                allow_undefined_type,
                original_ref,
            )

        # Get model properties
        # If this schema is referenced, use the ref name as the title if no title is provided
        if original_ref and "title" not in schema:
            ref_parts = original_ref.split("/")
            title = ref_parts[-1] if ref_parts else "DynamicModel"
        else:
            title = schema.get("title", "DynamicModel")
        description = schema.get("description")
        properties = schema.get("properties", {})
        required = schema.get("required", [])

        # Extract model-level json_schema_extra
        model_extra = {
            key: value
            for key, value in schema.items()
            if key not in self.STANDARD_MODEL_PROPERTIES
        }

        # Build field definitions
        fields = {}
        for field_name, field_schema in properties.items():
            field_type = self._get_field_type(
                field_schema,
                root_schema,
                allow_undefined_array_items,
                allow_undefined_type,
                populate_by_name,
            )
            model_field_name, alias = self._sanitize_field_name(field_name, set(properties))
            field_info = self._build_field_info(field_schema, field_name in required, alias=alias)
            fields[model_field_name] = (field_type, field_info)

        # Create the model with or without json_schema_extra
        if model_extra:
            # Create a dynamic base class with the config
            class DynamicBase(self.base_model_type):
                model_config = ConfigDict(
                    json_schema_extra=model_extra,
                    populate_by_name=populate_by_name,
                )

            model = create_model(title, __base__=DynamicBase, **fields)
        else:
            model = create_model(
                title,
                __base__=self.base_model_type,
                __config__=ConfigDict(populate_by_name=populate_by_name),
                **fields,
            )

        if description:
            model.__doc__ = description

        # Cache the model if it was referenced
        if original_ref:
            self._model_cache[original_ref] = model
            # Remove from building set
            self._building_models.discard(original_ref)
            # Mark model for rebuild after all models are created
            self._models_to_rebuild.add(model)

        # If this is the top-level call (no models being built), rebuild all models
        if not self._building_models and self._models_to_rebuild:
            # Build namespace with all models
            namespace = {m.__name__: m for m in self._models_to_rebuild}
            # Rebuild all models to resolve forward references
            for m in self._models_to_rebuild:
                m.model_rebuild(_types_namespace=namespace)
            # Clear the rebuild set
            self._models_to_rebuild.clear()

        return model

    def _create_array_root_model(
        self,
        schema: Dict[str, Any],
        root_schema: Dict[str, Any],
        allow_undefined_array_items: bool,
        original_ref: Optional[str],
    ) -> Type[T]:
        """
        Creates a RootModel for top-level array schemas.

        Args:
            schema: The array schema definition
            root_schema: The root schema containing definitions
            allow_undefined_array_items: Whether to allow arrays without items
            original_ref: The original reference if this was a $ref

        Returns:
            A RootModel class that validates arrays
        """
        # Get the title for the model
        if original_ref and "title" not in schema:
            ref_parts = original_ref.split("/")
            title = ref_parts[-1] if ref_parts else "DynamicModel"
        else:
            title = schema.get("title", "DynamicModel")

        # Get description
        description = schema.get("description")

        # Resolve the item type
        items_schema = schema.get("items")
        if not items_schema:
            if allow_undefined_array_items:
                item_type = Any
            else:
                from .exceptions import TypeError

                raise TypeError("Array type must specify 'items' schema")
        else:
            item_type = self._get_field_type(
                items_schema, root_schema, allow_undefined_array_items
            )

        # Determine if we need to use Set or List
        if schema.get("uniqueItems", False):
            array_type = Set[item_type]
        else:
            array_type = List[item_type]

        # Build constraints for the array
        constraints = self.constraint_builder.build_constraints(schema)

        # Extract model-level json_schema_extra (non-standard properties)
        model_extra = {
            key: value
            for key, value in schema.items()
            if key not in self.STANDARD_MODEL_PROPERTIES
        }

        # Create the RootModel class dynamically
        # RootModel requires the type to be specified as a generic parameter
        # We create a class that properly inherits from RootModel[array_type]

        # Build the class namespace
        namespace = {}
        if description:
            namespace["__doc__"] = description

        # Add model_config if we have extra properties
        if model_extra:
            namespace["model_config"] = ConfigDict(json_schema_extra=model_extra)

        # Apply constraints to the root field using Annotated
        if constraints:
            # Use Annotated to add Field constraints to the array type
            root_type = Annotated[array_type, Field(**constraints)]
        else:
            root_type = array_type
        namespace["__annotations__"] = {"root": root_type}

        # Create the RootModel subclass
        model = type(title, (RootModel[array_type],), namespace)

        # Cache the model if it was referenced
        if original_ref:
            self._model_cache[original_ref] = model
            # Mark model for rebuild if needed
            self._models_to_rebuild.add(model)

        # Rebuild models if this is the top-level call
        if not self._building_models and self._models_to_rebuild:
            namespace = {m.__name__: m for m in self._models_to_rebuild}
            for m in self._models_to_rebuild:
                m.model_rebuild(_types_namespace=namespace)
            self._models_to_rebuild.clear()

        return model

    def _create_scalar_root_model(
        self,
        schema: Dict[str, Any],
        root_schema: Dict[str, Any],
        allow_undefined_array_items: bool,
        allow_undefined_type: bool,
        original_ref: Optional[str],
    ) -> Type[T]:
        """
        Creates a RootModel for top-level scalar schemas (string, integer, etc).

        Args:
            schema: The JSON Schema fragment describing the scalar value.
            root_schema: The root JSON Schema document containing shared definitions.
            allow_undefined_array_items: Whether to allow array items without an explicit type
                when resolving nested schemas.
            allow_undefined_type: Whether to allow schemas that do not declare an explicit
                ``type`` field when resolving the scalar type.
            original_ref: The original JSON Schema ``$ref`` string for this schema, if any.

        Returns:
            Type[T]: A Pydantic ``RootModel`` subclass that represents the scalar schema.
        """
        # Get the title for the model
        if original_ref and "title" not in schema:
            ref_parts = original_ref.split("/")
            title = ref_parts[-1] if ref_parts else "DynamicModel"
        else:
            title = schema.get("title", "DynamicModel")

        description = schema.get("description")

        # Resolve the scalar type using existing logic
        scalar_type = self.type_resolver.resolve_type(
            schema, root_schema, allow_undefined_array_items, allow_undefined_type
        )

        # Build constraints
        constraints = self.constraint_builder.build_constraints(schema)

        # Extract model-level json_schema_extra
        model_extra = {
            key: value
            for key, value in schema.items()
            if key not in self.STANDARD_MODEL_PROPERTIES
        }

        # Build the class namespace
        namespace = {}
        if description:
            namespace["__doc__"] = description

        if model_extra:
            namespace["model_config"] = ConfigDict(json_schema_extra=model_extra)

        # Apply constraints using Annotated if needed
        if isinstance(constraints, dict) and constraints:
            root_type = Annotated[scalar_type, Field(**constraints)]
        else:
            root_type = scalar_type

        namespace["__annotations__"] = {"root": root_type}

        model = type(title, (RootModel[scalar_type],), namespace)

        if original_ref:
            self._model_cache[original_ref] = model
            self._models_to_rebuild.add(model)

        if not self._building_models and self._models_to_rebuild:
            ns = {m.__name__: m for m in self._models_to_rebuild}
            for m in self._models_to_rebuild:
                m.model_rebuild(_types_namespace=ns)
            self._models_to_rebuild.clear()

        return model

    def _get_field_type(
        self,
        field_schema: Dict[str, Any],
        root_schema: Dict[str, Any],
        allow_undefined_array_items: bool = False,
        allow_undefined_type: bool = False,
        populate_by_name: bool = False,
    ) -> Any:
        """Resolves the Python type for a field schema."""
        # Store the original ref if present
        original_ref = field_schema.get("$ref")

        if "$ref" in field_schema:
            if original_ref in self._ref_type_cache:
                return self._ref_type_cache[original_ref]

            # Check if this reference is already being built (recursive reference)
            if original_ref in self._building_models:
                # Return cached model if available
                if original_ref in self._model_cache:
                    return self._model_cache[original_ref]
                # Otherwise, return a forward reference (string)
                # Extract model name from ref
                ref_parts = original_ref.split("/")
                model_name = ref_parts[-1] if ref_parts else "DynamicModel"
                return model_name

            # Check if we've already built this model
            if original_ref in self._model_cache:
                return self._model_cache[original_ref]

            # Mark this ref as being built before resolving
            self._building_models.add(original_ref)

            field_schema = self.reference_resolver.resolve_ref(
                field_schema["$ref"], field_schema, root_schema
            )

        # Handle combiners
        if "allOf" in field_schema:
            return self.combiner_handler.handle_all_of(
                field_schema["allOf"],
                root_schema,
                allow_undefined_array_items,
                allow_undefined_type,
            )
        if "anyOf" in field_schema:
            return self.combiner_handler.handle_any_of(
                field_schema["anyOf"],
                root_schema,
                allow_undefined_array_items,
                allow_undefined_type,
            )
        if "oneOf" in field_schema:
            return self.combiner_handler.handle_one_of(
                field_schema["oneOf"],
                root_schema,
                allow_undefined_array_items,
                allow_undefined_type,
            )

        # Handle arrays by recursively processing items
        if field_schema.get("type") == "array":
            items_schema = field_schema.get("items")
            if not items_schema:
                if allow_undefined_array_items:
                    return List[Any]
                else:
                    from .exceptions import TypeError

                    raise TypeError("Array type must specify 'items' schema")

            # Recursively process the items schema through the model builder
            # This ensures that object types get proper models created
            item_type = self._get_field_type(
                items_schema,
                root_schema,
                allow_undefined_array_items,
                allow_undefined_type,
                populate_by_name,
            )

            if field_schema.get("uniqueItems", False):
                return Set[item_type]
            return List[item_type]

        # Handle nested objects by recursively creating models
        if field_schema.get("type") == "object" and "properties" in field_schema:
            # Build the model and cache/cleanup if we tracked this ref
            # Pass the original ref so the model can be named correctly
            model = self.create_pydantic_model(
                field_schema,
                root_schema,
                allow_undefined_array_items,
                allow_undefined_type,
                populate_by_name,
                _schema_ref=original_ref,
            )

            # If we were tracking a ref for this field, cache and cleanup
            if original_ref and original_ref in self._building_models:
                self._model_cache[original_ref] = model
                self._building_models.discard(original_ref)
                # Mark model for rebuild after all models are created
                self._models_to_rebuild.add(model)

            return model

        return self.type_resolver.resolve_type(
            schema=field_schema,
            root_schema=root_schema,
            allow_undefined_array_items=allow_undefined_array_items,
            allow_undefined_type=allow_undefined_type,
        )

    def _build_field_info(
        self,
        field_schema: Dict[str, Any],
        required: bool,
        alias: Optional[str] = None,
    ) -> Field:
        """Creates a Pydantic Field with constraints from schema.
        If the field_name is invalid as a model field name, it adds the original name as an alias.
        """
        field_kwargs = {}

        # Add constraints
        constraints = self.constraint_builder.build_constraints(field_schema)
        if isinstance(constraints, type):
            pass  # Type will be handled by type_resolver
        elif isinstance(constraints, dict):
            field_kwargs.update(constraints)

        # Handle description
        if "description" in field_schema:
            field_kwargs["description"] = field_schema["description"]

        # Handle default value
        if "default" in field_schema:
            field_kwargs["default"] = field_schema["default"]
        elif not required:
            field_kwargs["default"] = None

        # Handle alias
        if alias:
            field_kwargs["alias"] = alias

        # Extract field-level json_schema_extra
        field_extra = {
            key: value
            for key, value in field_schema.items()
            if key not in self.STANDARD_FIELD_PROPERTIES
        }

        if field_extra:
            field_kwargs["json_schema_extra"] = field_extra

        return Field(**field_kwargs)

    @staticmethod
    def _sanitize_field_name(field_name: str, invalid_names: set) -> tuple[str, Optional[str]]:
        """Sanitizes field names to be valid Pydantic field names.
        Returns the sanitized field name and an optional alias if the name was changed.
        """
        # If the sanitized name is different, return the alias
        sanitized_name = field_name
        alias = None
        if field_name.startswith("_"):
            sanitized_name, alias = field_name[1:], field_name
            if sanitized_name in invalid_names:
                raise SchemaError(
                    f"Duplicate field name after sanitization: '{sanitized_name}'\nPydantic does "
                    "not support field names starting with underscores when another field would "
                    "result in the same name."
                )
        return sanitized_name, alias
