"""Raw record loading and the FieldSpec-based default conversion.

Supported sources:
- local file:  .jsonl / .json / .csv / .tsv / .parquet
- local dir:   looks for ``{subset}_{split}.jsonl`` etc. (and csv/tsv/parquet)
- hub id:      HuggingFace ``datasets`` (optional dependency, imported on demand)
- modelscope:  ``params={'hub': 'modelscope'}`` — native ModelScope raw-file
  HTTP download (no ``modelscope`` package needed, immune to its datasets
  version pinning). File selection mirrors the local-dir convention and
  understands HF-style shards (``test-00000-of-00001.parquet``).
"""

import csv
import json
import os
import sys
import time
import re
import shutil
import urllib.request
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Union

from .sample import Sample
from .spec import DatasetSpec, FieldSpec

_SUPPORTED_EXTS = ['.jsonl', '.json', '.csv', '.tsv', '.parquet']
# params keys consumed by the loader itself; never forwarded to load_dataset()
_RESERVED_PARAMS = {'hub', 'ms_files', 'hf_files', 'filter_column'}
_MS_API = 'https://www.modelscope.cn/api/v1/datasets'


def get_cache_root() -> Path:
    """Canonical evalharness cache root ($EVALHARNESS_CACHE or ~/.cache/evalharness)."""
    return Path(os.environ.get('EVALHARNESS_CACHE', '~/.cache/evalharness')).expanduser()


def set_cache_root(path) -> None:
    """Override the cache root at runtime (also moves the .raw blob store)."""
    os.environ['EVALHARNESS_CACHE'] = str(path)


def load_raw_records(spec: DatasetSpec, raw_dir: Optional[Path] = None) -> List[Dict[str, Any]]:
    """Load raw records (list of dicts) from the spec's source.

    When ``raw_dir`` is given, a copy of the *native* source data is placed
    there (downloaded blob / copy of the local file / record dump for hub
    sources) so every cache entry is self-contained: raw in, samples out.
    """
    source = spec.source
    if os.path.exists(source):
        records, native_files = _load_local(source, spec)
        if raw_dir is not None and native_files:
            _link_or_copy_all(native_files, raw_dir)
    elif spec.params.get('hub') == 'modelscope':
        records = _load_from_modelscope(spec, raw_dir)
    elif spec.params.get('hub') == 'hf_raw':
        records = _load_from_hf_raw(spec, raw_dir)
    else:
        records = _load_from_hub(spec)
        if raw_dir is not None:
            # hub-native packaging lives in the HF cache; keep an exact,
            # pre-conversion record dump so the entry is self-contained
            raw_dir.mkdir(parents=True, exist_ok=True)
            with open(raw_dir / 'records.jsonl', 'w', encoding='utf-8') as f:
                for r in records:
                    f.write(json.dumps(r, ensure_ascii=False, default=str) + '\n')
    records = _apply_subset_filter(spec, records)
    return records


def _load_local(path: str, spec: DatasetSpec) -> tuple:
    """Load from a local path; also return the native file(s) to preserve."""
    if os.path.isfile(path):
        return _read_file(path), [Path(path)]
    # directory: follow the <subset>_<split>.<ext> / <subset>.<ext> convention
    for ext in _SUPPORTED_EXTS:
        for candidate in (f'{spec.subset}_{spec.split}{ext}', f'{spec.subset}{ext}', f'{spec.split}{ext}'):
            full = os.path.join(path, candidate)
            if os.path.exists(full):
                return _read_file(full), [Path(full)]
    expected = [os.path.join(path, f'{spec.subset}_{spec.split}{e}') for e in _SUPPORTED_EXTS]
    available = sorted(f for f in os.listdir(path) if os.path.splitext(f)[1] in _SUPPORTED_EXTS)
    raise FileNotFoundError(
        f'no dataset file found for subset={spec.subset!r} split={spec.split!r} in {path!r}.\n'
        f'Expected one of:\n  - ' + '\n  - '.join(expected) + '\n'
        f'Available: {available or "(none)"}'
    )


def _link_or_copy_all(files: List[Path], raw_dir: Path) -> None:
    """Place native files in raw_dir (hardlink when possible, else copy)."""
    raw_dir.mkdir(parents=True, exist_ok=True)
    for src in files:
        dst = raw_dir / src.name
        if dst.exists() or src.parent.resolve() == raw_dir.resolve():
            continue
        try:
            os.link(src, dst)
        except OSError:
            shutil.copy2(src, dst)


def _read_file(path: str) -> List[Dict[str, Any]]:
    ext = os.path.splitext(path)[1]
    if ext == '.parquet':
        try:
            import pyarrow.parquet as pq
        except ImportError:
            raise ImportError(
                f'{path} is parquet; install the reader first: pip install "evalharness[all]"'
            )
        return pq.read_table(path).to_pylist()
    with open(path, encoding='utf-8') as f:
        if ext == '.jsonl':
            return [json.loads(line) for line in f if line.strip()]
        if ext == '.json':
            data = json.load(f)
            return data if isinstance(data, list) else [data]
        if ext in ('.csv', '.tsv'):
            return list(csv.DictReader(f, delimiter='\t' if ext == '.tsv' else ','))
    raise ValueError(f'unsupported file format: {path}')


def _apply_subset_filter(spec: DatasetSpec, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """When ``params['filter_column']`` is set, subset selects that column's value.

    Used when a mirror packs all subsets into one file with a category column
    (e.g. evalscope/cmmlu: one parquet, ``category`` = subject).
    """
    col = spec.params.get('filter_column')
    if not col or spec.subset in ('default', 'all'):
        return records
    return [r for r in records if r.get(col) == spec.subset]


# ---------------- ModelScope raw-file loading (no modelscope package) ----------------


def _ms_api_json(url: str) -> Dict[str, Any]:
    req = urllib.request.Request(url, headers={'User-Agent': 'evalharness/0.1'})
    with urllib.request.urlopen(req, timeout=60) as resp:
        return json.loads(resp.read().decode('utf-8'))


def _ms_list_files(repo: str, root: str = '', depth: int = 0) -> List[str]:
    """Recursively list blob paths under a ModelScope dataset repo directory."""
    url = f'{_MS_API}/{repo}/repo/tree?Revision=master' + (f'&Root={root}' if root else '')
    try:
        data = _ms_api_json(url).get('Data') or {}
    except Exception:
        return []
    files: List[str] = []
    for entry in data.get('Files') or []:
        if entry.get('Type') == 'blob':
            files.append(entry['Path'])
        elif entry.get('Type') == 'tree' and depth < 4:
            files.extend(_ms_list_files(repo, entry['Path'], depth + 1))
    return files


def _ms_match_files(spec: DatasetSpec, files: List[str]) -> List[str]:
    """Pick repo files for this subset/split, mirroring the local-dir convention.

    Priority: explicit ``params['ms_files']`` > exact name match > HF-style
    shard match (``{split}-00000-of-00001`` / ``{subset}_{split}-...`` /
    ``{subset}_0``) > single-data-file repos. When the subset names a parent
    directory, matches inside it win (e.g. livecodebench ``release_latest/``).
    """
    explicit = spec.params.get('ms_files')
    if explicit:
        # paths support {subset}/{split} templating (e.g. tau2 domains)
        return [p.format(subset=spec.subset, split=spec.split) for p in explicit]

    def match_one(path: str) -> bool:
        base = os.path.basename(path)
        name = os.path.splitext(base)[0]
        for ext in _SUPPORTED_EXTS:
            if not base.endswith(ext):
                continue
            stem = base[: -len(ext)]
            if stem in (f'{spec.subset}_{spec.split}', spec.subset, spec.split):
                return True
            if stem == f'{spec.subset}{spec.split}':
                return True
            shard = (
                # tolerate a content-hash suffix after the shard pattern
                # (e.g. test-00000-of-00001-6282153cc50a2626.parquet)
                rf'{re.escape(spec.split)}-\d+-of-\d+(-[0-9a-f]+)?$'
                rf'|{re.escape(spec.subset)}_{re.escape(spec.split)}[-_].*'
                rf'|{re.escape(spec.subset)}[-_]\d+$'
            )
            if re.fullmatch(shard, stem, flags=re.ASCII):
                return True
        return False

    candidates = [f for f in files if match_one(f)]
    if not candidates:
        data_files = [f for f in files if os.path.splitext(f)[1] in _SUPPORTED_EXTS]
        if len(data_files) == 1:
            return data_files
        return []
    in_subset_dir = [f for f in candidates if os.path.dirname(f) == spec.subset]
    return sorted(in_subset_dir or candidates)


def _parquet_ok(path: Path) -> bool:
    """Cheap integrity check: parquet files end with magic 'PAR1'."""
    try:
        with open(path, 'rb') as f:
            f.seek(-4, 2)
            return f.read(4) == b'PAR1'
    except OSError:
        return False



def _download_with_progress(resp, out, filename: str) -> None:
    """Stream a download with a rich byte-level progress bar.

    Used for the minute-scale dataset fetches (ModelScope blobs, HF raw
    files); silent (plain streaming) when rich is unavailable or output
    is redirected."""
    total = 0
    try:
        total = int(resp.headers.get('Content-Length') or 0)
    except (TypeError, ValueError):
        total = 0
    if not sys.stdout.isatty():
        while True:
            chunk = resp.read(1 << 20)
            if not chunk:
                break
            out.write(chunk)
        return
    try:
        from rich.progress import (BarColumn, DownloadColumn, Progress,
                                   SpinnerColumn, TextColumn,
                                   TimeElapsedColumn, TransferSpeedColumn)

        pg = Progress(SpinnerColumn(),
                      TextColumn('[cyan]{task.fields[name]}[/cyan]'),
                      BarColumn(), DownloadColumn(), TransferSpeedColumn(),
                      TextColumn('•'), TimeElapsedColumn(),
                      transient=True)
        with pg:
            t = pg.add_task('download', total=total or None, name=filename)
            while True:
                chunk = resp.read(1 << 20)
                if not chunk:
                    break
                out.write(chunk)
                pg.advance(t, len(chunk))
        return
    except ImportError:
        pass
    while True:
        chunk = resp.read(1 << 20)
        if not chunk:
            break
        out.write(chunk)


def _ms_download(repo: str, path: str, dest_dir: Path) -> Path:
    """Download one repo file into the raw cache (content-addressed, reused)."""
    dest = dest_dir / os.path.basename(path)
    if dest.exists() and dest.stat().st_size > 0 \
            and (os.path.splitext(dest)[1] != '.parquet' or _parquet_ok(dest)):
        return dest
    if dest.exists():  # truncated/corrupt (e.g. an interrupted download)
        dest.unlink()
    dest_dir.mkdir(parents=True, exist_ok=True)
    url = f'{_MS_API}/{repo}/repo?Revision=master&FilePath={path}'
    tmp = dest.with_name(dest.name + f'.part-{os.getpid()}')
    req = urllib.request.Request(url, headers={'User-Agent': 'evalharness/0.1'})
    with urllib.request.urlopen(req, timeout=600) as resp, open(tmp, 'wb') as out:
        _download_with_progress(resp, out, dest.name)
    os.replace(tmp, dest)
    return dest


def _load_from_modelscope(spec: DatasetSpec, raw_dir: Optional[Path] = None) -> List[Dict[str, Any]]:
    import hashlib

    files = _ms_list_files(spec.source)
    if not files:
        raise FileNotFoundError(f'no files found on ModelScope dataset {spec.source!r}')
    selected = _ms_match_files(spec, files)
    if not selected:
        raise FileNotFoundError(
            f'no data file for subset={spec.subset!r} split={spec.split!r} in '
            f'modelscope {spec.source!r}. Available (first 10): {files[:10]}'
        )
    # shared blob store: download once per repo, hardlink into each cache entry
    blob_dir = get_cache_root() / '.raw' / hashlib.md5(spec.source.encode()).hexdigest()[:10]
    records: List[Dict[str, Any]] = []
    blobs: List[Path] = []
    for path in selected:
        blobs.append(_ms_download(spec.source, path, blob_dir))
        records.extend(_read_file(str(blobs[-1])))
    if raw_dir is not None:
        _link_or_copy_all(blobs, raw_dir)
    return records


# ---------------- HuggingFace raw-file loading (no auth, exact bytes) ----------------
# For repos whose layout datasets.load_dataset cannot express (raw json trees,
# shard dirs without configs) or that only exist as plain files. Honors
# $HF_ENDPOINT (e.g. https://hf-mirror.com).


def _hf_base() -> str:
    return os.environ.get('HF_ENDPOINT', 'https://huggingface.co').rstrip('/')


def _hf_list_files(repo: str, root: str = '', depth: int = 0) -> List[str]:
    """Recursively list file paths under an HF dataset repo directory."""
    url = f'{_hf_base()}/api/datasets/{repo}/tree/main' + (f'/{root}' if root else '') + '?limit=1000'
    req = urllib.request.Request(url, headers={'User-Agent': 'evalharness/0.1'})
    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            entries = json.loads(resp.read().decode('utf-8'))
    except Exception:
        return []
    files: List[str] = []
    for entry in entries:
        if entry.get('type') == 'file':
            files.append(entry['path'])
        elif entry.get('type') == 'directory' and depth < 4:
            files.extend(_hf_list_files(repo, entry['path'], depth + 1))
    return files


def _hf_match_files(spec: DatasetSpec, files: List[str]) -> List[str]:
    """Pick repo files for this subset/split (see module docstring conventions)."""
    explicit = spec.params.get('hf_files')
    if explicit:
        return [p.format(subset=spec.subset, split=spec.split) for p in explicit]

    def ext_ok(p: str) -> bool:
        return os.path.splitext(p)[1] in _SUPPORTED_EXTS

    data = [f for f in files if ext_ok(f)]
    # multi-window subset, e.g. 'release_v2_v4' = union of release_v2..v4
    # shards (LCB release-window convention: per-release files are disjoint)
    m = re.fullmatch(r'(release_v\d+)((?:_v\d+)+)', spec.subset)
    if m:
        members = [m.group(1)] + [f'release{x}' for x in m.group(2).split('_') if x]
        picked = []
        for mem in members:
            picked += [
                f for f in data
                if re.fullmatch(rf'{re.escape(mem)}[-_]\d+(-of-\d+)?'
                                rf'|{re.escape(mem)}_{re.escape(spec.split)}[-_].*'
                                rf'|{re.escape(mem)}',
                                os.path.splitext(os.path.basename(f))[0])
            ]
        if picked:
            # deterministic: member order, shard order within each member
            return [f for mem in members for f in sorted(
                p for p in picked
                if os.path.splitext(os.path.basename(p))[0].startswith(mem))]
    if spec.subset != 'default':
        # subset wins exclusively: never ALSO match bare-split shards in the
        # same dir (repos like sam-paech LCB mix test-*.parquet and
        # release_v*-*.parquet in one data/ folder)
        in_dir = [f for f in data if os.path.dirname(f) == spec.subset]
        if in_dir:
            return sorted(in_dir)
        exact = []
        for f in data:
            stem = os.path.splitext(os.path.basename(f))[0]
            if stem in (f'{spec.subset}_{spec.split}', spec.subset):
                exact.append(f)
        if exact:
            return sorted(exact)
        shards = [
            f for f in data
            if re.fullmatch(rf'{re.escape(spec.subset)}[-_]\d+(-of-\d+)?'
                            rf'|{re.escape(spec.subset)}_{re.escape(spec.split)}[-_].*',
                            os.path.splitext(os.path.basename(f))[0])
        ]
        if shards:
            return sorted(shards)
    else:
        exact = [
            f for f in data
            if os.path.splitext(os.path.basename(f))[0] == spec.split
        ]
        if exact:
            return sorted(exact)
        shards = [
            f for f in data
            if re.fullmatch(rf'{re.escape(spec.split)}[-_]\d+(-of-\d+)?',
                            os.path.splitext(os.path.basename(f))[0])
        ]
        if shards:
            return sorted(shards)
    if len(data) == 1:
        return data
    return []


def _hf_download(repo: str, path: str, dest_dir: Path) -> Path:
    """Download one repo file (follows the CDN redirect) into the blob store."""
    dest = dest_dir / os.path.basename(path)
    if dest.exists() and dest.stat().st_size > 0 \
            and (os.path.splitext(dest)[1] != '.parquet' or _parquet_ok(dest)):
        return dest
    if dest.exists():  # truncated/corrupt (e.g. an interrupted download)
        dest.unlink()
    dest_dir.mkdir(parents=True, exist_ok=True)
    url = f'{_hf_base()}/datasets/{repo}/resolve/main/{path}'
    tmp = dest.with_name(dest.name + f'.part-{os.getpid()}')
    req = urllib.request.Request(url, headers={'User-Agent': 'evalharness/0.1'})
    with urllib.request.urlopen(req, timeout=1800) as resp, open(tmp, 'wb') as out:
        _download_with_progress(resp, out, dest.name)
    os.replace(tmp, dest)
    return dest


def _load_from_hf_raw(spec: DatasetSpec, raw_dir: Optional[Path] = None) -> List[Dict[str, Any]]:
    import hashlib

    # the mirror's tree API degrades at night and returns TRUNCATED listings
    # (a repo with 9 parquets listed as [.gitattributes, README.md]); the
    # blobs from a previous run are already in the shared store -- use them
    blob_dir = get_cache_root() / '.raw' / hashlib.md5(spec.source.encode()).hexdigest()[:10]

    def _cached_blobs() -> List[str]:
        if not blob_dir.exists():
            return []
        have = [f.name for f in blob_dir.iterdir()
                if os.path.splitext(f)[1] in _SUPPORTED_EXTS]
        pref = [f for f in have if f.startswith(spec.subset)] or have
        return pref

    files = _hf_list_files(spec.source)
    selected = _hf_match_files(spec, files) if files else []
    if not selected:
        cached = _cached_blobs()
        if cached:
            print(f'· listing unavailable/degraded -- using {len(cached)} '
                  f'cached file(s) from {blob_dir}', flush=True)
            records: List[Dict[str, Any]] = []
            for f in sorted(cached):
                records.extend(_read_file(str(blob_dir / f)))
            if raw_dir is not None:
                _link_or_copy_all([blob_dir / f for f in sorted(cached)], raw_dir)
            return records
        raise FileNotFoundError(
            f'no data file for subset={spec.subset!r} split={spec.split!r} in '
            f'HF {spec.source!r}. Available (first 10): {files[:10]}'
        )
    blob_dir = get_cache_root() / '.raw' / hashlib.md5(spec.source.encode()).hexdigest()[:10]
    records: List[Dict[str, Any]] = []
    blobs: List[Path] = []
    for path in selected:
        blobs.append(_hf_download(spec.source, path, blob_dir))
        records.extend(_read_file(str(blobs[-1])))
    if raw_dir is not None:
        _link_or_copy_all(blobs, raw_dir)
    return records


def _load_from_hub(spec: DatasetSpec) -> List[Dict[str, Any]]:
    try:
        import datasets
    except ImportError:
        raise ImportError(
            f'dataset {spec.name!r} lives on a hub ({spec.source!r}); '
            'install it first: pip install evalharness (light deps are default)'
        )
    _probe_hub_reachable(spec)
    kwargs = {k: v for k, v in spec.params.items() if k not in _RESERVED_PARAMS}
    # filter_column subsets select rows by column value post-load, so the hub
    # load itself must always use the default config.
    subset = None if (spec.subset == 'default' or spec.params.get('filter_column')) else spec.subset
    print(f'· downloading {spec.source} from HF hub ({_hf_base()}) ...', flush=True)
    t0 = time.monotonic()
    ds = datasets.load_dataset(spec.source, subset, split=spec.split, revision=spec.version, **kwargs)
    print(f'· downloaded + parsed {len(ds)} records in {time.monotonic() - t0:.0f}s', flush=True)
    return [dict(r) for r in ds]


_HF_MIRROR_FALLBACKS = ['https://hf-mirror.com']


def _probe_hub_reachable(spec: DatasetSpec) -> None:
    """Fail fast when the HF endpoint is unreachable -- otherwise the hub
    client retries silently for minutes and looks like a hang.

    When the DEFAULT endpoint is down but a mirror answers, switch to it
    automatically (one-line notice, same behavior as --hf-endpoint)."""
    base = _hf_base()
    err = None
    for candidate in [base] + (_HF_MIRROR_FALLBACKS if base == 'https://huggingface.co' else []):
        try:
            req = urllib.request.Request(candidate, method='HEAD',
                                         headers={'User-Agent': 'evalharness/0.1'})
            urllib.request.urlopen(req, timeout=8)
            if candidate != base:
                os.environ['HF_ENDPOINT'] = candidate
                print(f'· {base} unreachable -- falling back to {candidate}',
                      flush=True)
            return
        except Exception as e:
            err = e
    raise RuntimeError(
        f'HuggingFace endpoint unreachable: {base}'
        + (f' (mirrors tried: {", ".join(_HF_MIRROR_FALLBACKS)})'
           if base == 'https://huggingface.co' else '')
        + f'\n  reason: {type(err).__name__}: {str(err)[:80]}\n'
        f'  dataset {spec.source!r} cannot download. Fix:\n'
        f'  evalharness eval run ... --hf-endpoint <a reachable endpoint>') from err


def field_spec_to_record_fn(fields: FieldSpec) -> Callable[[Dict[str, Any]], Sample]:
    """Build the default record->Sample converter from a FieldSpec."""

    def convert(record: Dict[str, Any]) -> Sample:
        target = record.get(fields.target, '')
        if isinstance(target, (int, float)):
            target = str(target)
        metadata = {k: record.get(k) for k in fields.metadata}
        return Sample(
            input=record.get(fields.input, ''),
            choices=record.get(fields.choices) if fields.choices in record else None,
            target=target,
            id=record.get(fields.id) if fields.id else None,
            metadata=metadata,
        )

    return convert