"""Hub/local dataset fetcher with checksum, TTL, and deterministic cache index."""

from __future__ import annotations

import hashlib
import json
import os
import time
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Optional, Tuple
from urllib.parse import urlparse

from model_osmosis.services.data_registry import DataSourceProfile


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 _utc_now() -> str:
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


def _cache_index_path() -> Path:
    return expand_data_home() / "index.json"


def _cache_file_path(profile: DataSourceProfile) -> Path:
    parsed = urlparse(profile.uri)
    ext = Path(parsed.path).suffix or ".bin"
    digest = hashlib.sha256(f"{profile.dataset_id}:{profile.revision}:{profile.uri}".encode("utf-8")).hexdigest()[:24]
    return expand_data_home() / f"{digest}{ext}"


def _load_index() -> Dict[str, object]:
    path = _cache_index_path()
    if not path.exists():
        return {"schema_version": 1, "updated_at": "", "entries": {}}
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return {"schema_version": 1, "updated_at": "", "entries": {}}
    if not isinstance(payload, dict):
        return {"schema_version": 1, "updated_at": "", "entries": {}}
    entries = payload.get("entries", {})
    if not isinstance(entries, dict):
        entries = {}
    return {
        "schema_version": 1,
        "updated_at": str(payload.get("updated_at", "")),
        "entries": entries,
    }


def _save_index(index_payload: Dict[str, object]) -> None:
    path = _cache_index_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    normalized = dict(index_payload)
    normalized["schema_version"] = 1
    normalized["updated_at"] = _utc_now()
    path.write_text(json.dumps(normalized, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def _entry_key(profile: DataSourceProfile) -> str:
    return f"{profile.dataset_id}:{profile.revision}"


def _is_fresh(entry: Dict[str, object], *, ttl_hours: float) -> bool:
    fetched_at = str(entry.get("fetched_at", ""))
    if not fetched_at:
        return False
    try:
        ts = time.strptime(fetched_at, "%Y-%m-%dT%H:%M:%SZ")
        fetched_epoch = time.mktime(ts)
    except Exception:
        return False
    age_seconds = max(0.0, time.time() - float(fetched_epoch))
    return age_seconds <= float(ttl_hours) * 3600.0


def _download_with_limit(uri: str, *, target: Path, max_download_bytes: int, timeout_seconds: int) -> Tuple[int, str]:
    sha256 = hashlib.sha256()
    total = 0
    chunk_size = 1024 * 64
    with urllib.request.urlopen(uri, timeout=timeout_seconds) as resp:  # nosec B310
        with open(target, "wb") as fh:
            while True:
                chunk = resp.read(chunk_size)
                if not chunk:
                    break
                total += len(chunk)
                if total > max_download_bytes:
                    raise RuntimeError(
                        f"Download exceeded max size ({max_download_bytes} bytes) for uri={uri}"
                    )
                fh.write(chunk)
                sha256.update(chunk)
    return total, sha256.hexdigest()


@dataclass(frozen=True)
class FetchResult:
    path: Path
    cache_hit: bool
    downloaded_bytes: int
    checksum_sha256: str
    source: str
    reason: str

    def to_dict(self) -> Dict[str, object]:
        return {
            "path": str(self.path),
            "cache_hit": self.cache_hit,
            "downloaded_bytes": int(self.downloaded_bytes),
            "checksum_sha256": self.checksum_sha256,
            "source": self.source,
            "reason": self.reason,
        }


def fetch_profile_data(
    profile: DataSourceProfile,
    *,
    max_download_gb: float = 1.0,
    ttl_hours: float = 24.0 * 7.0,
    offline: bool = False,
    timeout_seconds: int = 20,
) -> FetchResult:
    """Fetch dataset profile to local cache and return cache metadata."""
    home = expand_data_home()
    home.mkdir(parents=True, exist_ok=True)
    cache_path = _cache_file_path(profile)
    index = _load_index()
    entries = index.get("entries", {})
    if not isinstance(entries, dict):
        entries = {}
    key = _entry_key(profile)
    existing = entries.get(key, {})
    if isinstance(existing, dict):
        if cache_path.exists() and _is_fresh(existing, ttl_hours=float(ttl_hours)):
            return FetchResult(
                path=cache_path,
                cache_hit=True,
                downloaded_bytes=0,
                checksum_sha256=str(existing.get("checksum_sha256", "")),
                source="cache",
                reason="fresh_cache",
            )

    if offline:
        if cache_path.exists():
            checksum = hashlib.sha256(cache_path.read_bytes()).hexdigest()
            return FetchResult(
                path=cache_path,
                cache_hit=True,
                downloaded_bytes=0,
                checksum_sha256=checksum,
                source="cache",
                reason="offline_cache_fallback",
            )
        raise RuntimeError(
            f"Offline mode enabled and no cached data found for dataset_id={profile.dataset_id}"
        )

    max_bytes = int(max(0.001, float(max_download_gb)) * (1024**3))
    downloaded_bytes, checksum = _download_with_limit(
        profile.uri,
        target=cache_path,
        max_download_bytes=max_bytes,
        timeout_seconds=int(max(1, int(timeout_seconds))),
    )
    expected = str(profile.checksum_sha256 or "").strip().lower()
    if expected and checksum.lower() != expected:
        try:
            cache_path.unlink(missing_ok=True)  # type: ignore[arg-type]
        except Exception:
            pass
        raise RuntimeError(
            f"Checksum mismatch for dataset_id={profile.dataset_id}: expected={expected} got={checksum}"
        )

    entries[key] = {
        "dataset_id": profile.dataset_id,
        "revision": profile.revision,
        "path": str(cache_path),
        "checksum_sha256": checksum,
        "size_bytes": int(downloaded_bytes),
        "license": profile.license,
        "hub": profile.hub,
        "fetched_at": _utc_now(),
    }
    index["entries"] = entries
    _save_index(index)
    return FetchResult(
        path=cache_path,
        cache_hit=False,
        downloaded_bytes=int(downloaded_bytes),
        checksum_sha256=checksum,
        source="hub",
        reason="downloaded",
    )
