from typing import (
    Any,
    Callable,
    Dict,
    List,
    Optional,
    Tuple,
    Type,
    Union,
    get_args,
    get_origin,
)

from pydantic import AliasChoices, AliasGenerator, AliasPath, BaseModel
from pydantic.fields import FieldInfo
from guardrails.classes.output_type import OutputTypes
from guardrails.classes.schema.processed_schema import ProcessedSchema
from guardrails_ai.types import Validator as ValidatorReference
from guardrails.logger import logger
from guardrails.types import (
    ModelOrListOfModels,
    ModelOrListOrDict,
    ModelOrModelUnion,
)
from guardrails.utils.safe_get import safe_get
from guardrails.utils.validator_utils import safe_get_validator
from guardrails.validator_base import Validator


def _resolve_alias(alias: Union[str, AliasPath, AliasChoices]) -> List[str]:
    aliases = []
    if isinstance(alias, str):
        aliases.append(alias)
    elif isinstance(alias, AliasPath):
        aliases.append(".".join(str(alias.path)))
    elif isinstance(alias, AliasChoices):
        for choice in alias.choices:
            aliases.extend(_resolve_alias(choice))

    return aliases


def _collect_aliases(
    field: Union[FieldInfo, AliasGenerator], field_name: str, model: Type[BaseModel]
) -> List[str]:
    aliases = []

    if field.alias:
        if isinstance(field.alias, str):
            aliases.append(field.alias)
        elif isinstance(field.alias, Callable):
            aliases.append(field.alias(field_name))

    if field.serialization_alias:
        if isinstance(field.serialization_alias, str):
            aliases.append(field.serialization_alias)
        elif isinstance(field.serialization_alias, Callable):
            aliases.append(field.serialization_alias(field_name))

    if field.validation_alias:
        if isinstance(field.validation_alias, Callable):
            aliases.extend(_resolve_alias(field.validation_alias(field_name)))
        else:
            aliases.extend(_resolve_alias(field.validation_alias))

    alias_generator = model.model_config.get("alias_generator")
    if alias_generator:
        if isinstance(alias_generator, Callable):
            aliases.append(alias_generator(field_name))
        elif isinstance(alias_generator, AliasGenerator):
            return _collect_aliases(alias_generator, field_name, model)

    return aliases


def is_base_model_type(any_type: Any) -> bool:
    try:
        inherits_from_base_model = issubclass(any_type, BaseModel)
        return inherits_from_base_model
    except TypeError:
        return False


def get_base_model(
    pydantic_class: ModelOrListOrDict,
) -> Tuple[ModelOrModelUnion, Any, Optional[Any]]:
    schema_model = pydantic_class
    type_origin = get_origin(pydantic_class)
    key_type_origin = None

    if type_origin is list:
        item_types = get_args(pydantic_class)
        if len(item_types) > 1:
            raise ValueError("List data type must have exactly one child.")
        item_type = safe_get(item_types, 0)
        if not item_type or not issubclass(item_type, BaseModel):
            raise ValueError("List item type must be a Pydantic model.")
        schema_model = item_type
    elif type_origin is dict:
        key_value_types = get_args(pydantic_class)
        value_type = safe_get(key_value_types, 1)
        key_type_origin = safe_get(key_value_types, 0)
        if not value_type or not issubclass(value_type, BaseModel):
            raise ValueError("Dict value type must be a Pydantic model.")
        schema_model = value_type
    elif type_origin is Union:
        union_members = get_args(pydantic_class)
        model_members = list(filter(is_base_model_type, union_members))
        if len(model_members) > 0:
            schema_model = Union[tuple(union_members)]  # type: ignore
            return (schema_model, type_origin, key_type_origin)

    if not is_base_model_type(schema_model):
        raise ValueError(
            "'output_class' must be of Type[pydantic.BaseModel]"
            " or List[Type[pydantic.BaseModel]]!"
        )

    return (schema_model, type_origin, key_type_origin)


def try_get_base_model(
    pydantic_class: ModelOrListOrDict,
) -> Tuple[Optional[Type[BaseModel]], Optional[Any], Optional[Any]]:
    try:
        model, type_origin, key_type_origin = get_base_model(pydantic_class)
        return (model, type_origin, key_type_origin)
    except ValueError:
        return (None, None, None)
    except TypeError:
        return (None, None, None)


def extract_union_member(
    member: Type,
    processed_schema: ProcessedSchema,
    json_path: str,
    aliases: List[str],
) -> Type:
    aliases = aliases or []
    field_model, field_type_origin, key_type_origin = try_get_base_model(member)
    if not field_model:
        return member
    if field_type_origin is Union:
        union_members = get_args(field_model)
        extracted_union_members = []
        for m in union_members:
            extracted_union_members.append(
                extract_union_member(m, processed_schema, json_path, aliases)
            )
        return Union[tuple(extracted_union_members)]  # type: ignore

    else:
        extracted_field_model = extract_validators(
            model=field_model,
            processed_schema=processed_schema,
            json_path=json_path,
            aliases=aliases,
        )
        if field_type_origin is list:
            return List[extracted_field_model]
        elif field_type_origin is dict:
            return Dict[key_type_origin, extracted_field_model]  # type: ignore
        return extracted_field_model


def extract_validators(
    model: Type[BaseModel],
    processed_schema: ProcessedSchema,
    json_path: str,
    aliases: Optional[List[str]] = None,
) -> Type[BaseModel]:
    aliases = aliases or []
    for field_name in model.model_fields:
        alias_paths = []
        field_path = f"{json_path}.{field_name}"
        # alias_paths.append(field_path)
        for alias_path in aliases:
            alias_paths.append(f"{alias_path}.{field_name}")
        field: FieldInfo = model.model_fields[field_name]
        for alias in _collect_aliases(field, field_name, model):
            alias_paths.append(f"{json_path}.{alias}")
            for alias_path in aliases:
                alias_paths.append(f"{alias_path}.{alias}")

        if field.json_schema_extra is not None and isinstance(
            field.json_schema_extra, dict
        ):
            # NOTE: It's impossible to copy a class type so using
            #   'pop' here mutates the original Pydantic Model.
            # Using 'get' adds a pointless 'validators' field to the
            #   json schema but that doesn't break anything.
            validators = field.json_schema_extra.get("validators", [])

            if not isinstance(validators, list) and not isinstance(
                validators, Validator
            ):
                logger.warning(
                    f"Invalid value assigned to {field_name}.validators! {validators}"
                )
                continue
            validator_instances: List[Validator] = []

            # Only for backwards compatibility
            if isinstance(validators, Validator):
                validator_instances.append(validators)
            else:
                validator_list = [
                    safe_get_validator(v)  # type: ignore
                    for v in validators
                ]
                validator_instances.extend([v for v in validator_list if v is not None])
            all_paths = [field_path]
            all_paths.extend(alias_paths)
            for path in all_paths:
                entry = processed_schema.validator_map.get(path, [])
                entry.extend(validator_instances)
                processed_schema.validator_map[path] = entry
                validator_references = [
                    ValidatorReference(
                        id=v.rail_alias,
                        on=path,
                        on_fail=v.on_fail_descriptor,  # type: ignore
                        kwargs=v.get_args(),
                    )
                    for v in validator_instances
                ]
                processed_schema.validators.extend(validator_references)
        if field.annotation:
            field_model, field_type_origin, key_type_origin = try_get_base_model(
                field.annotation
            )
            if field_model:
                if field_type_origin is Union:
                    union_members = list(get_args(field_model))
                    extracted_union_members = []
                    for m in union_members:
                        extracted_union_members.append(
                            extract_union_member(
                                m,
                                processed_schema=processed_schema,
                                json_path=field_path,
                                aliases=alias_paths,
                            )
                        )

                    model.model_fields[field_name].annotation = Union[  # type: ignore
                        tuple(extracted_union_members)  # type: ignore
                    ]
                else:
                    extracted_field_model = extract_validators(
                        model=field_model,
                        processed_schema=processed_schema,
                        json_path=field_path,
                        aliases=alias_paths,
                    )
                    if field_type_origin is list:
                        model.model_fields[field_name].annotation = List[
                            extracted_field_model
                        ]
                    elif field_type_origin is dict:
                        model.model_fields[field_name].annotation = Dict[
                            key_type_origin, extracted_field_model  # type: ignore
                        ]
                    else:
                        model.model_fields[
                            field_name
                        ].annotation = extracted_field_model  # noqa
    return model


def pydantic_to_json_schema(
    pydantic_class: Type[BaseModel], type_origin: Optional[Any] = None
) -> Dict[str, Any]:
    # Convert Pydantic model to JSON schema
    json_schema = pydantic_class.model_json_schema()
    json_schema["title"] = pydantic_class.__name__

    if type_origin is list:
        json_schema = {
            "title": f"Array<{json_schema.get('title')}>",
            "type": "array",
            "items": json_schema,
        }

    return json_schema


def pydantic_model_to_schema(
    pydantic_class: ModelOrListOfModels,
) -> ProcessedSchema:
    processed_schema = ProcessedSchema(validators=[], validator_map={})

    schema_model, type_origin, _key_type_origin = get_base_model(pydantic_class)

    processed_schema.output_type = (
        OutputTypes.LIST if type_origin is list else OutputTypes.DICT
    )

    model = extract_validators(schema_model, processed_schema, "$")
    json_schema = pydantic_to_json_schema(model, type_origin)
    processed_schema.json_schema = json_schema

    return processed_schema