from __future__ import annotations

import dataclasses
import inspect
from typing import Any, get_args, get_origin, get_type_hints
from ._exceptions import SchemaGenerationError

def _normalize_annotation(annotation: Any) -> Any:
    origin = get_origin(annotation)
    args = get_args(annotation)

    if annotation is Any:
        return object

    if origin is None:
        if inspect.isclass(annotation) and dataclasses.is_dataclass(annotation):
            return generate_schema(annotation)
        return annotation

    if origin in {list, set, frozenset, tuple}:
        if not args:
            return [object]
        return [_normalize_annotation(args[0])]

    if origin is dict:
        if len(args) == 2:
            key_type, value_type = args
            if key_type not in (str, Any, object):
                raise SchemaGenerationError("Only dicts with string-like keys are supported")
            return {"<value>": _normalize_annotation(value_type)}
        return {"<value>": object}

    if str(origin).endswith("Union") or "types.UnionType" in str(type(origin)):
        normalized = tuple(_normalize_annotation(arg) for arg in args if arg is not type(None))
        return normalized[0] if len(normalized) == 1 else normalized

    return object


def _dataclass_schema(cls: type[Any]) -> dict[str, Any]:
    hints = get_type_hints(cls)
    return {field.name: _normalize_annotation(hints.get(field.name, Any)) for field in dataclasses.fields(cls)}


def _typing_hints_schema(cls: type[Any]) -> dict[str, Any]:
    return {name: _normalize_annotation(annotation) for name, annotation in get_type_hints(cls).items()}


def _pydantic_schema(model: type[Any]) -> dict[str, Any]:
    schema: dict[str, Any] = {}
    if hasattr(model, "model_fields"):
        for name, field in getattr(model, "model_fields").items():
            schema[name] = _normalize_annotation(getattr(field, "annotation", Any))
        return schema
    if hasattr(model, "__fields__"):
        for name, field in getattr(model, "__fields__").items():
            schema[name] = _normalize_annotation(getattr(field, "outer_type_", None) or getattr(field, "type_", Any))
        return schema
    raise SchemaGenerationError("Provided model is not a supported pydantic model")


def postprocess_generated_schema(schema: Any) -> Any:
    if isinstance(schema, dict):
        if "<value>" in schema and len(schema) == 1:
            return dict
        return {key: postprocess_generated_schema(value) for key, value in schema.items()}
    if isinstance(schema, list):
        return [postprocess_generated_schema(schema[0])] if schema else list
    if isinstance(schema, tuple):
        return tuple(postprocess_generated_schema(item) for item in schema)
    return schema


def generate_schema(target: type[Any], *, simplify_dict_values: bool = True) -> dict[str, Any]:
    if not inspect.isclass(target):
        raise TypeError("target must be a class")
    if dataclasses.is_dataclass(target):
        schema = _dataclass_schema(target)
    elif hasattr(target, "model_fields") or hasattr(target, "__fields__"):
        schema = _pydantic_schema(target)
    else:
        schema = _typing_hints_schema(target)
    if not schema:
        raise SchemaGenerationError(f"Could not infer schema from target: {target!r}")
    return postprocess_generated_schema(schema) if simplify_dict_values else schema