"""Data planning helpers for local expansion workflows."""

from __future__ import annotations

import hashlib
import os
import urllib.request
from pathlib import Path
from typing import Any, Dict, Iterable, Optional, Sequence, Tuple
from urllib.parse import urlparse

import numpy as np

from model_osmosis.services.data_fetcher import fetch_profile_data
from model_osmosis.services.data_registry import DataSourceProfile
from model_osmosis.services.data_resolver import resolve_data_source


SUPPORTED_EXPAND_TASKS = {
    "auto",
    "text",
    "vision",
    "audio",
    "code",
    "classification",
    "regression",
}


def expand_data_home() -> Path:
    env = str(os.environ.get("OSMOSIS_EXPAND_DATA_HOME", "")).strip()
    if env:
        return Path(env)
    return Path(".model_osmosis") / "expand_data"


def _is_uri(value: str) -> bool:
    parsed = urlparse(str(value))
    return parsed.scheme in {"http", "https", "file"}


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


def infer_expand_task(*, task: Optional[str], model_ref: str, model_obj: Any) -> str:
    """Infer task from model reference/name when task is `auto`."""
    normalized = _normalize_task(task)
    if normalized != "auto":
        return normalized

    ref = str(model_ref or "").lower()
    hints = " ".join(
        [
            ref,
            str(type(model_obj).__name__).lower(),
            str(type(model_obj).__module__).lower(),
        ]
    )
    if any(token in hints for token in ("whisper", "wav2vec", "audio", "speech")):
        return "audio"
    if any(token in hints for token in ("vit", "resnet", "clip", "vision", "image")):
        return "vision"
    if ref.startswith("hf:"):
        return "text"
    if any(token in hints for token in ("llama", "gpt", "mistral", "qwen", "text", "token")):
        return "text"
    if "code" in hints:
        return "code"
    if "regress" in hints:
        return "regression"
    return "classification"


def _download_cached(uri: str) -> Tuple[Path, bool]:
    home = expand_data_home()
    home.mkdir(parents=True, exist_ok=True)
    parsed = urlparse(uri)
    ext = Path(parsed.path).suffix or ".bin"
    digest = hashlib.sha256(uri.encode("utf-8")).hexdigest()[:16]
    target = home / f"{digest}{ext}"
    if target.exists():
        return target, True
    with urllib.request.urlopen(uri, timeout=20) as resp:  # nosec B310
        data = resp.read()
    target.write_bytes(data)
    return target, False


def _parse_allowed_licenses(values: Optional[Sequence[str]]) -> Tuple[str, ...]:
    if values:
        out = []
        seen = set()
        for raw in values:
            token = str(raw or "").strip().lower()
            if not token or token in seen:
                continue
            seen.add(token)
            out.append(token)
        return tuple(out)
    env_raw = str(os.environ.get("OSMOSIS_EXPAND_ALLOWED_LICENSES", "")).strip()
    if not env_raw:
        return ()
    return tuple(
        token.strip().lower()
        for token in env_raw.split(",")
        if token.strip()
    )


def _load_numpy_dataset(path: Path) -> np.ndarray:
    loaded = np.load(path, allow_pickle=True)
    if isinstance(loaded, np.ndarray):
        arr = loaded
    else:
        files = list(getattr(loaded, "files", []) or [])
        if not files:
            raise RuntimeError(f"Dataset file has no arrays: {path}")
        arr = loaded[files[0]]
    arr = np.asarray(arr)
    if arr.ndim == 1:
        arr = arr.reshape(-1, 1)
    return arr.astype("float32", copy=False)


def _synthetic_data(*, seed: int, feature_dim: int = 8, samples: int = 256) -> np.ndarray:
    rng = np.random.default_rng(int(seed))
    return rng.standard_normal((int(samples), int(feature_dim))).astype("float32")


def _to_task_batch(
    *,
    task: str,
    seed: int,
    array_data: np.ndarray,
    feature_dim: int,
) -> Any:
    """Convert numeric dataset into task-shaped runtime sample batch."""
    normalized = _normalize_task(task)
    arr = np.asarray(array_data)
    if arr.ndim == 1:
        arr = arr.reshape(-1, 1)

    if normalized in {"classification", "regression", "auto"}:
        return arr.astype("float32", copy=False)

    if normalized in {"text", "code"}:
        if arr.ndim > 2:
            arr = arr.reshape(arr.shape[0], -1)
        token_ids = np.abs(np.rint(arr * 1000.0)).astype("int64", copy=False) % 32000
        if token_ids.ndim == 1:
            token_ids = token_ids.reshape(1, -1)
        if token_ids.shape[1] < 8:
            repeats = int(np.ceil(8.0 / float(max(1, token_ids.shape[1]))))
            token_ids = np.tile(token_ids, (1, repeats))
        token_ids = token_ids[:, : min(256, token_ids.shape[1])]
        return {
            "kind": "text",
            "input_ids": token_ids,
            "attention_mask": np.ones_like(token_ids, dtype=np.int64),
        }

    if normalized == "vision":
        if arr.ndim == 4 and arr.shape[1] in {1, 3}:
            images = arr.astype("float32", copy=False)
        else:
            batch = int(max(1, arr.shape[0]))
            flat = arr.reshape(batch, -1)
            images = np.resize(flat, (batch, 3, 32, 32)).astype("float32", copy=False)
        return {"kind": "vision", "images": images}

    if normalized == "audio":
        if arr.ndim == 2:
            waveforms = arr.astype("float32", copy=False)
        else:
            batch = int(max(1, arr.shape[0]))
            flat = arr.reshape(batch, -1)
            waveforms = np.resize(flat, (batch, 16000)).astype("float32", copy=False)
        return {"kind": "audio", "waveforms": waveforms, "sample_rate": 16000}

    # Defensive fallback for future tasks.
    rng = np.random.default_rng(int(seed))
    fallback = rng.standard_normal((256, int(max(1, feature_dim)))).astype("float32")
    return fallback


def _batch_report_fields(batch: Any) -> Dict[str, Any]:
    if isinstance(batch, np.ndarray):
        return {
            "task_batch_kind": "tabular",
            "task_batch_type": "ndarray",
            "task_batch_shape": [int(x) for x in batch.shape],
        }
    if isinstance(batch, dict):
        kind = str(batch.get("kind", "unknown"))
        shape: Any = []
        if kind in {"text", "code"}:
            ids = np.asarray(batch.get("input_ids"))
            shape = [int(x) for x in ids.shape]
        elif kind == "vision":
            images = np.asarray(batch.get("images"))
            shape = [int(x) for x in images.shape]
        elif kind == "audio":
            wav = np.asarray(batch.get("waveforms"))
            shape = [int(x) for x in wav.shape]
        return {
            "task_batch_kind": kind,
            "task_batch_type": "dict",
            "task_batch_shape": shape,
        }
    return {
        "task_batch_kind": "unknown",
        "task_batch_type": str(type(batch).__name__),
        "task_batch_shape": [],
    }


def plan_expand_data(
    *,
    task: str,
    dataset: Optional[str],
    seed: int,
    feature_dim: int = 8,
    model_ref: str = "",
    data_mode: str = "auto",
    max_download_gb: float = 1.0,
    data_ttl_hours: float = 24.0 * 7.0,
    allow_licenses: Optional[Sequence[str]] = None,
    offline: bool = False,
    registry: Optional[Iterable[DataSourceProfile]] = None,
) -> Tuple[np.ndarray, Dict[str, Any]]:
    """Return expansion sample data and deterministic planner report."""
    normalized_task = _normalize_task(task)
    normalized_mode = str(data_mode or "auto").strip().lower()
    if normalized_mode not in {"auto", "off"}:
        raise ValueError("Unsupported data_mode. Expected one of: auto, off")
    allowed_licenses = _parse_allowed_licenses(allow_licenses)

    report: Dict[str, Any] = {
        "task": normalized_task,
        "data_mode": normalized_mode,
        "dataset_source": "synthetic",
        "cache_hit": False,
        "cache_path": "",
        "dataset_input": str(dataset or ""),
        "model_ref": str(model_ref or ""),
        "selection_reason": "",
        "dataset_id": "",
        "dataset_revision": "",
        "source_hub": "",
        "download_mb": 0.0,
        "license": "",
        "fallback_reason": "",
        "allowed_licenses": list(allowed_licenses),
        "license_policy": "any" if not allowed_licenses else "allowlist",
        "planner_route": "synthetic",
        "route_reason": "",
        "task_batch_kind": "",
        "task_batch_type": "",
        "task_batch_shape": [],
    }

    if dataset:
        dataset_s = str(dataset).strip()
        if _is_uri(dataset_s):
            cached, hit = _download_cached(dataset_s)
            report.update(
                {
                    "dataset_source": "user_uri",
                    "cache_hit": bool(hit),
                    "cache_path": str(cached),
                }
            )
            arr = _load_numpy_dataset(cached)
            batch = _to_task_batch(
                task=normalized_task,
                seed=int(seed),
                array_data=arr,
                feature_dim=int(max(1, feature_dim)),
            )
            report.update(_batch_report_fields(batch))
            report["sample_shape"] = list(arr.shape)
            report["planner_route"] = "dataset:user_uri"
            report["route_reason"] = "explicit_dataset_uri"
            return batch, report

        dataset_path = Path(dataset_s)
        if not dataset_path.exists():
            raise FileNotFoundError(f"Dataset not found: {dataset_path}")
        arr = _load_numpy_dataset(dataset_path)
        batch = _to_task_batch(
            task=normalized_task,
            seed=int(seed),
            array_data=arr,
            feature_dim=int(max(1, feature_dim)),
        )
        report.update(
            {
                "dataset_source": "user_path",
                "cache_hit": True,
                "cache_path": str(dataset_path.resolve()),
                "sample_shape": list(arr.shape),
                "planner_route": "dataset:user_path",
                "route_reason": "explicit_dataset_path",
            }
        )
        report.update(_batch_report_fields(batch))
        return batch, report

    if normalized_mode == "off":
        arr = _synthetic_data(seed=seed, feature_dim=max(1, int(feature_dim)))
        batch = _to_task_batch(
            task=normalized_task,
            seed=int(seed),
            array_data=arr,
            feature_dim=int(max(1, feature_dim)),
        )
        report["dataset_source"] = "synthetic_data_mode_off"
        report["sample_shape"] = list(arr.shape)
        report["planner_route"] = "synthetic"
        report["route_reason"] = "data_mode_off"
        report.update(_batch_report_fields(batch))
        return batch, report

    try:
        resolution = resolve_data_source(
            task=normalized_task,
            model_ref=str(model_ref or ""),
            seed=int(seed),
            allowed_licenses=allowed_licenses,
            registry=registry,
        )
        report["selection_reason"] = str(resolution.selection_reason)
        report["considered_sources"] = list(resolution.considered)

        profile = resolution.selected
        if profile is not None:
            fetched = fetch_profile_data(
                profile,
                max_download_gb=float(max_download_gb),
                ttl_hours=float(data_ttl_hours),
                offline=bool(offline),
                timeout_seconds=20,
            )
            arr = _load_numpy_dataset(fetched.path)
            batch = _to_task_batch(
                task=normalized_task,
                seed=int(seed),
                array_data=arr,
                feature_dim=int(max(1, feature_dim)),
            )
            report.update(
                {
                    "dataset_source": "auto_hub",
                    "cache_hit": bool(fetched.cache_hit),
                    "cache_path": str(fetched.path),
                    "sample_shape": list(arr.shape),
                    "dataset_id": profile.dataset_id,
                    "dataset_revision": profile.revision,
                    "source_hub": profile.hub,
                    "download_mb": round(float(fetched.downloaded_bytes) / (1024.0 * 1024.0), 4),
                    "license": profile.license,
                    "planner_route": f"auto:{profile.hub}",
                    "route_reason": str(resolution.selection_reason or "best_score"),
                }
            )
            report.update(_batch_report_fields(batch))
            return batch, report

        # Backward-compatible single URI fallback.
        open_hub_uri = str(os.environ.get("OSMOSIS_EXPAND_OPEN_HUB_URI", "")).strip()
        if open_hub_uri:
            cached, hit = _download_cached(open_hub_uri)
            arr = _load_numpy_dataset(cached)
            batch = _to_task_batch(
                task=normalized_task,
                seed=int(seed),
                array_data=arr,
                feature_dim=int(max(1, feature_dim)),
            )
            report.update(
                {
                    "dataset_source": "open_hub_env_fallback",
                    "cache_hit": bool(hit),
                    "cache_path": str(cached),
                    "sample_shape": list(arr.shape),
                    "source_hub": "env",
                    "selection_reason": "open_hub_env_fallback",
                    "planner_route": "auto:env_fallback",
                    "route_reason": "open_hub_env_fallback",
                }
            )
            report.update(_batch_report_fields(batch))
            return batch, report
    except Exception as exc:
        report["fallback_reason"] = f"auto_data_failed:{type(exc).__name__}"
        report["auto_data_error"] = str(exc)

    arr = _synthetic_data(seed=seed, feature_dim=max(1, int(feature_dim)))
    batch = _to_task_batch(
        task=normalized_task,
        seed=int(seed),
        array_data=arr,
        feature_dim=int(max(1, feature_dim)),
    )
    report["sample_shape"] = list(arr.shape)
    if not report.get("selection_reason"):
        report["selection_reason"] = "synthetic_fallback"
    report["planner_route"] = "synthetic"
    report["route_reason"] = "synthetic_fallback"
    report.update(_batch_report_fields(batch))
    return batch, report
