"""Deterministic registry of expansion calibration/training data sources."""

from __future__ import annotations

import os
from dataclasses import dataclass, asdict
from typing import Dict, Iterable, List, Optional


SUPPORTED_DATA_TASKS = {
    "text",
    "vision",
    "audio",
    "code",
    "classification",
    "regression",
}


@dataclass(frozen=True)
class DataSourceProfile:
    """Single dataset source candidate for expansion planner."""

    dataset_id: str
    task: str
    uri: str
    hub: str = "unknown"
    revision: str = "latest"
    license: str = "unknown"
    priority: int = 100
    min_quality: float = 0.0
    checksum_sha256: str = ""
    max_rows_hint: int = 0

    def to_dict(self) -> Dict[str, object]:
        return asdict(self)


def _normalize_task(task: str) -> str:
    value = str(task or "").strip().lower()
    if value not in SUPPORTED_DATA_TASKS:
        raise ValueError(
            f"Unsupported data task '{value}'. Expected one of: {sorted(SUPPORTED_DATA_TASKS)}"
        )
    return value


def _parse_env_sources(task: str) -> List[DataSourceProfile]:
    """Parse env override list: OSMOSIS_EXPAND_DATA_<TASK>=uri1,uri2."""
    env_key = f"OSMOSIS_EXPAND_DATA_{task.upper()}"
    raw = str(os.environ.get(env_key, "")).strip()
    if not raw:
        return []
    out: List[DataSourceProfile] = []
    for idx, token in enumerate(raw.split(","), start=1):
        uri = token.strip()
        if not uri:
            continue
        out.append(
            DataSourceProfile(
                dataset_id=f"{task}_env_{idx}",
                task=task,
                uri=uri,
                hub="env",
                revision="env",
                license="unknown",
                priority=10 + idx,
                min_quality=0.1,
            )
        )
    return out


def _builtin_sources(task: str) -> List[DataSourceProfile]:
    """Built-in placeholder hub sources (safe fallback if unavailable)."""
    # These are intentionally best-effort sources; planner will gracefully
    # fallback if unreachable.
    if task == "text":
        return [
            DataSourceProfile(
                dataset_id="hf_text_calibration_small",
                task=task,
                uri="https://huggingface.co/datasets/tiiuae/falcon-refinedweb/resolve/main/README.md",
                hub="huggingface",
                revision="main",
                license="apache-2.0",
                priority=40,
                min_quality=0.30,
            )
        ]
    if task == "vision":
        return [
            DataSourceProfile(
                dataset_id="hf_vision_calibration_small",
                task=task,
                uri="https://huggingface.co/datasets/beans/resolve/main/README.md",
                hub="huggingface",
                revision="main",
                license="mit",
                priority=40,
                min_quality=0.28,
            )
        ]
    if task == "audio":
        return [
            DataSourceProfile(
                dataset_id="hf_audio_calibration_small",
                task=task,
                uri="https://huggingface.co/datasets/librispeech_asr/resolve/main/README.md",
                hub="huggingface",
                revision="main",
                license="cc-by-4.0",
                priority=45,
                min_quality=0.26,
            )
        ]
    if task == "code":
        return [
            DataSourceProfile(
                dataset_id="hf_code_calibration_small",
                task=task,
                uri="https://huggingface.co/datasets/code_search_net/resolve/main/README.md",
                hub="huggingface",
                revision="main",
                license="mit",
                priority=40,
                min_quality=0.32,
            )
        ]
    if task == "classification":
        return [
            DataSourceProfile(
                dataset_id="hf_classification_calibration_small",
                task=task,
                uri="https://huggingface.co/datasets/ag_news/resolve/main/README.md",
                hub="huggingface",
                revision="main",
                license="cc-by-4.0",
                priority=45,
                min_quality=0.24,
            )
        ]
    if task == "regression":
        return [
            DataSourceProfile(
                dataset_id="hf_regression_calibration_small",
                task=task,
                uri="https://huggingface.co/datasets/scikit-learn/diabetes/resolve/main/README.md",
                hub="huggingface",
                revision="main",
                license="bsd-3-clause",
                priority=50,
                min_quality=0.22,
            )
        ]
    return []


def load_data_registry(*, task: str, include_env: bool = True) -> List[DataSourceProfile]:
    """Load data source candidates for a specific task."""
    normalized_task = _normalize_task(task)
    items: List[DataSourceProfile] = []
    if include_env:
        items.extend(_parse_env_sources(normalized_task))
    items.extend(_builtin_sources(normalized_task))
    # Ensure deterministic ordering.
    return sorted(items, key=lambda item: (int(item.priority), str(item.dataset_id), str(item.uri)))


def filter_allowed_licenses(
    profiles: Iterable[DataSourceProfile],
    allowed_licenses: Optional[Iterable[str]],
) -> List[DataSourceProfile]:
    if not allowed_licenses:
        return list(profiles)
    allowed = {str(x).strip().lower() for x in allowed_licenses if str(x).strip()}
    if not allowed:
        return list(profiles)
    out = []
    for profile in profiles:
        if str(profile.license).strip().lower() in allowed:
            out.append(profile)
    return out
