from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Any, Mapping


VALUE_TYPE_MAP = {
    "string": "str",
    "large_string": "str",
    "int16": "int",
    "int32": "int",
    "int64": "int",
    "uint16": "int",
    "uint32": "int",
    "uint64": "int",
    "float16": "float",
    "float32": "float",
    "float64": "float",
    "bool": "bool",
}


@dataclass
class FieldSpec:
    name: str
    annotation: str


class SchemaRenderer:
    """Render Python source for a dataset schema dataclass."""

    def __init__(self, dataset_id: str, config_name: str | None):
        self.dataset_id = dataset_id
        self.config_name = config_name

    def render(self, dataset_info: Mapping[str, Any]) -> str:
        features = dataset_info.get("features") or {}
        field_specs: list[FieldSpec] = []
        typing_imports: set[str] = {"ClassVar", "Optional"}

        for name, feature in features.items():
            annotation, extra_imports = self._annotation_for_feature(feature)
            typing_imports.update(extra_imports)
            field_specs.append(FieldSpec(name=name, annotation=annotation))

        lines = ["from dataclasses import dataclass"]
        if typing_imports:
            lines.append("from typing import " + ", ".join(sorted(typing_imports)))
        lines.append("")
        lines.append("from schemas_hf.base import DatasetSchema")
        lines.append("from schemas_hf.registry import register_schema")
        lines.append("")
        lines.append("")
        lines.append("@register_schema")
        lines.append("@dataclass")
        lines.append(f"class {self._class_name()}(DatasetSchema):")
        lines.append(f"    DATASET_NAME: ClassVar[str] = \"{self.dataset_id}\"")
        config_value = f"\"{self.config_name}\"" if self.config_name else "None"
        lines.append(f"    CONFIG_NAME: ClassVar[Optional[str]] = {config_value}")

        if field_specs:
            for field in field_specs:
                lines.append(f"    {field.name}: {field.annotation}")
        else:
            lines.append("    pass")

        lines.append("")
        return "\n".join(lines)

    def _class_name(self) -> str:
        parts = self.dataset_id.split("/")
        name_parts = [_pascal_case(part) for part in parts]
        if self.config_name:
            name_parts.append(_pascal_case(self.config_name))
        return "".join(name_parts) + "Record"

    def _annotation_for_feature(self, feature: Any) -> tuple[str, set[str]]:
        if isinstance(feature, list):
            if not feature:
                return "list[Any]", {"Any"}
            inner_annotation, extra_imports = self._annotation_for_feature(feature[0])
            return f"list[{inner_annotation}]", extra_imports

        if not isinstance(feature, Mapping):
            return "Any", {"Any"}

        feature_type = feature.get("_type")
        if feature_type == "Value":
            dtype = feature.get("dtype")
            return VALUE_TYPE_MAP.get(dtype, "Any"), ({"Any"} if dtype not in VALUE_TYPE_MAP else set())

        if feature_type == "ClassLabel":
            num_classes = feature.get("num_classes") or len(feature.get("names") or [])
            if num_classes and num_classes <= 10:
                literal_values = ", ".join(str(i) for i in range(num_classes))
                return f"Literal[{literal_values}]", {"Literal"}
            return "int", set()

        if feature_type == "Sequence" and "feature" in feature:
            inner_annotation, extra_imports = self._annotation_for_feature(feature["feature"])
            return f"list[{inner_annotation}]", extra_imports

        if feature_type is None and any(isinstance(v, Mapping) for v in feature.values()):
            return "dict[str, Any]", {"Any"}

        return "Any", {"Any"}


def _pascal_case(value: str) -> str:
    tokens = re.split(r"[^a-zA-Z0-9]+", value)
    return "".join(token.capitalize() for token in tokens if token)
