from __future__ import annotations

from collections.abc import Iterable, Mapping
from pathlib import Path
from typing import Any, Dict

from .codegen import SchemaRenderer
from .naming import module_filename
from .sync import AUTO_HEADER


def infer_features(examples: Iterable[Mapping[str, Any]]) -> Dict[str, Dict[str, Any]]:
    """
    Infer a datasets-style features mapping from a collection of example rows.

    The output is compatible with SchemaRenderer, i.e. a mapping from column name
    to a feature dict such as:
      {"_type": "Value", "dtype": "int64"}
      {"_type": "Sequence", "feature": {"_type": "Value", "dtype": "string"}}
    """
    collected: list[Mapping[str, Any]] = list(examples)
    if not collected:
        return {}

    columns: dict[str, list[Any]] = {}
    for row in collected:
        for key, value in row.items():
            columns.setdefault(key, []).append(value)

    features: dict[str, dict[str, Any]] = {}
    for name, values in columns.items():
        features[name] = _infer_feature_for_values(values)
    return features


def write_schema_module(
    dataset_id: str,
    config_name: str | None,
    features: Mapping[str, Any],
    output_dir: Path,
) -> Path:
    """
    Materialize a schema module for a dataset given an inferred features mapping.

    Returns the path to the written module.
    """
    output_dir = Path(output_dir)
    module_path = output_dir / module_filename(dataset_id, config_name)
    renderer = SchemaRenderer(dataset_id, config_name)
    contents = AUTO_HEADER + renderer.render({"features": dict(features)})
    module_path.parent.mkdir(parents=True, exist_ok=True)
    module_path.write_text(contents)
    return module_path


def generate_schema_from_paths(
    dataset_id: str,
    paths: Iterable[str | Path],
    *,
    format: str,
    output_dir: Path,
    sample_rows: int = 1000,
) -> Path:
    """
    Convenience helper: sample examples from local files, infer features, and
    materialize a schema module.

    Currently supports:
      - format='jsonl'  (one JSON object per line)
    """
    format_normalized = format.lower()
    path_objs = [Path(p) for p in paths]
    if not path_objs:
        raise ValueError("No input paths provided.")

    if format_normalized == "jsonl":
        examples = _sample_jsonl_examples(path_objs, sample_rows)
    else:
        raise ValueError(
            f"Unsupported format: {format!r}. Supported formats: 'jsonl'."
        )

    if not examples:
        raise ValueError("No examples could be read from the provided paths.")

    features = infer_features(examples)
    return write_schema_module(dataset_id, None, features, output_dir)


def _sample_jsonl_examples(paths: list[Path], sample_rows: int) -> list[Dict[str, Any]]:
    import json

    examples: list[Dict[str, Any]] = []
    remaining = max(sample_rows, 0)
    for path in paths:
        if remaining <= 0:
            break
        if not path.exists() or not path.is_file():
            continue
        with path.open("r", encoding="utf-8") as f:
            for line in f:
                if remaining <= 0:
                    break
                line = line.strip()
                if not line:
                    continue
                try:
                    obj = json.loads(line)
                except json.JSONDecodeError:
                    continue
                if isinstance(obj, Mapping):
                    examples.append(dict(obj))
                    remaining -= 1
    return examples


def _infer_feature_for_values(values: list[Any]) -> Dict[str, Any]:
    """Infer a single feature spec from observed values for one column."""
    # Drop missing values from inference; if everything is missing, fall back to unknown.
    present = [v for v in values if v is not None]
    if not present:
        return {"_type": "Value", "dtype": "unknown"}

    kinds: set[str] = set()
    for v in present:
        if isinstance(v, bool):
            kinds.add("bool")
        elif isinstance(v, int):
            kinds.add("int")
        elif isinstance(v, float):
            kinds.add("float")
        elif isinstance(v, str):
            kinds.add("string")
        elif isinstance(v, (list, tuple)):
            kinds.add("list")
        elif isinstance(v, Mapping):
            kinds.add("dict")
        else:
            kinds.add("other")

    # Sequence types
    if kinds == {"list"}:
        return _infer_sequence_feature(present)  # type: ignore[arg-type]

    # Any mixture that includes dict/other is treated as opaque.
    if "dict" in kinds or "other" in kinds:
        return {"_type": "Value", "dtype": "unknown"}

    # Scalar numeric / string / bool cases
    if kinds == {"bool"}:
        return {"_type": "Value", "dtype": "bool"}
    if kinds == {"int"}:
        return {"_type": "Value", "dtype": "int64"}
    if kinds <= {"int", "float"}:
        return {"_type": "Value", "dtype": "float64"}
    if kinds == {"string"}:
        return {"_type": "Value", "dtype": "string"}

    # Mixed scalars (e.g. int + string) – fall back to unknown.
    return {"_type": "Value", "dtype": "unknown"}


def _infer_sequence_feature(values: list[list[Any]]) -> Dict[str, Any]:
    """
    Infer a Sequence feature from a collection of list values.

    Supports:
      - list[scalar]
      - list[list[scalar]]  (e.g. 2D grids)
    Anything more complex falls back to Sequence[Any].
    """
    # Collect inner elements
    inner_scalars: list[Any] = []
    nested_lists: list[list[Any]] = []
    for v in values:
        if not isinstance(v, (list, tuple)):
            continue
        for item in v:
            if isinstance(item, (list, tuple)):
                nested_lists.append(list(item))
            else:
                inner_scalars.append(item)

    # list[list[scalar]] case (e.g. ARCAGI style grids)
    if nested_lists and not inner_scalars:
        deepest_scalars: list[Any] = []
        for lst in nested_lists:
            for item in lst:
                if item is not None:
                    deepest_scalars.append(item)
        if not deepest_scalars:
            inner_feature = {"_type": "Value", "dtype": "unknown"}
        else:
            inner_feature = _infer_feature_for_values(deepest_scalars)
        return {
            "_type": "Sequence",
            "feature": {
                "_type": "Sequence",
                "feature": inner_feature,
            },
        }

    # list[scalar] case
    if inner_scalars and not nested_lists:
        inner_feature = _infer_feature_for_values(inner_scalars)
        return {"_type": "Sequence", "feature": inner_feature}

    # Mixed or empty – default to Sequence[Any]
    return {
        "_type": "Sequence",
        "feature": {"_type": "Value", "dtype": "unknown"},
    }


__all__ = ["infer_features", "write_schema_module", "generate_schema_from_paths"]

