"""Dataset card parser — extract provenance, licensing, and metadata from ML datasets.

Parses local dataset card files:
- HuggingFace ``dataset_info.json`` (auto-generated by 🤗 Datasets)
- HuggingFace dataset ``README.md`` YAML frontmatter (manual card)
- DVC ``.dvc`` tracking files (data versioning provenance)

Flags security/compliance concerns:
- ``UNLICENSED_DATASET`` — no license declaration (MEDIUM)
- ``NO_DATASET_CARD`` — README exists but lacks YAML frontmatter (LOW)
- ``UNVERSIONED_DATA`` — DVC file without md5/sha256 hash (LOW)
- ``REMOTE_DATA_SOURCE`` — DVC tracks external storage (INFO)
"""

from __future__ import annotations

import json
import logging
import re
from dataclasses import dataclass, field
from pathlib import Path

from agent_bom.parsers.compliance_tags import tag_dataset

logger = logging.getLogger(__name__)

# Skip directories during discovery
_SKIP_DIRS = {".git", "node_modules", "__pycache__", ".venv", "venv", ".tox", ".eggs"}

# YAML frontmatter regex (same pattern as skills.py)
_FRONTMATTER_RE = re.compile(r"\A---\s*\n(.*?)\n---\s*\n", re.DOTALL)

# Simple YAML key: value parser (avoids PyYAML dependency)
_YAML_KV_RE = re.compile(r"^(\w[\w_-]*)\s*:\s*(.+)$", re.MULTILINE)
_YAML_LIST_ITEM_RE = re.compile(r"^\s*-\s+(.+)$", re.MULTILINE)


@dataclass
class DatasetInfo:
    """Parsed dataset card metadata."""

    name: str = ""
    description: str = ""
    license: str = ""
    source_url: str = ""
    version: str = ""
    features: list[str] = field(default_factory=list)
    splits: dict[str, int] = field(default_factory=dict)
    size_bytes: int = 0
    citation: str = ""
    task_categories: list[str] = field(default_factory=list)
    languages: list[str] = field(default_factory=list)
    source_file: str = ""
    dvc_remote: str = ""
    dvc_md5: str = ""
    security_flags: list[dict] = field(default_factory=list)
    compliance_tags: dict[str, list[str]] = field(default_factory=dict)

    def to_dict(self) -> dict[str, object]:
        """Serialize to plain dict for JSON output."""
        d: dict[str, object] = {
            "name": self.name,
            "description": self.description[:300] if self.description else "",
            "license": self.license,
            "source_file": self.source_file,
        }
        if self.source_url:
            d["source_url"] = self.source_url
        if self.version:
            d["version"] = self.version
        if self.features:
            d["features"] = self.features
        if self.splits:
            d["splits"] = self.splits
        if self.size_bytes:
            d["size_bytes"] = self.size_bytes
        if self.citation:
            d["citation"] = self.citation[:500]
        if self.task_categories:
            d["task_categories"] = self.task_categories
        if self.languages:
            d["languages"] = self.languages
        if self.dvc_remote:
            d["dvc_remote"] = self.dvc_remote
        if self.dvc_md5:
            d["dvc_md5"] = self.dvc_md5
        if self.security_flags:
            d["security_flags"] = self.security_flags
        if self.compliance_tags:
            d["compliance_tags"] = self.compliance_tags
        return d


@dataclass
class DatasetScanResult:
    """Aggregated dataset scan results."""

    datasets: list[DatasetInfo] = field(default_factory=list)
    source_files: list[str] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "datasets": [d.to_dict() for d in self.datasets],
            "source_files": self.source_files,
            "warnings": self.warnings,
            "total_datasets": len(self.datasets),
            "flagged_count": sum(1 for d in self.datasets if d.security_flags),
        }


def _parse_simple_yaml(text: str) -> dict[str, str]:
    """Extract key-value pairs from simple YAML text (no nested structures)."""
    return dict(_YAML_KV_RE.findall(text))


def _parse_yaml_list(text: str, key: str) -> list[str]:
    """Extract a YAML list value for a given key."""
    # Find the key, then collect indented list items until next key
    pattern = re.compile(rf"^{re.escape(key)}\s*:\s*\n((?:\s+-\s+.+\n?)+)", re.MULTILINE)
    match = pattern.search(text)
    if not match:
        return []
    return _YAML_LIST_ITEM_RE.findall(match.group(1))


# ─── Parsers ────────────────────────────────────────────────────────────────


def parse_dataset_info_json(path: Path) -> DatasetInfo | None:
    """Parse a HuggingFace ``dataset_info.json`` file.

    This file is auto-generated by the 🤗 Datasets library and contains
    structured metadata about dataset features, splits, and size.
    """
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError) as exc:
        logger.warning("Failed to parse %s: %s", path, exc)
        return None

    if not isinstance(data, dict):
        return None

    info = DatasetInfo(
        name=data.get("dataset_name", data.get("config_name", path.parent.name)) or "",
        description=(data.get("description", "") or "")[:300],
        version=str(data.get("version", {}).get("version_str", ""))
        if isinstance(data.get("version"), dict)
        else str(data.get("version", "")),
        citation=data.get("citation", "") or "",
        source_file=str(path),
    )

    # Extract license
    lic = data.get("license", "")
    if isinstance(lic, list):
        info.license = lic[0] if lic else ""
    elif isinstance(lic, str):
        info.license = lic

    # Extract features (column names)
    features = data.get("features", {})
    if isinstance(features, dict):
        info.features = list(features.keys())[:50]  # cap at 50

    # Extract splits
    splits = data.get("splits", {})
    if isinstance(splits, dict):
        for split_name, split_data in splits.items():
            if isinstance(split_data, dict):
                info.splits[split_name] = split_data.get("num_examples", 0)
            elif isinstance(split_data, int):
                info.splits[split_name] = split_data

    # Extract size
    info.size_bytes = data.get("download_size", 0) or data.get("dataset_size", 0) or 0

    # Security flags
    if not info.license:
        info.security_flags.append(
            {
                "severity": "MEDIUM",
                "type": "UNLICENSED_DATASET",
                "description": (
                    f"Dataset '{info.name}' has no license declaration."
                    " Training on unlicensed data may violate IP rights (EU AI Act Art. 10)."
                ),
            }
        )

    return info


def parse_dataset_readme(path: Path) -> DatasetInfo | None:
    """Parse a HuggingFace dataset README.md with YAML frontmatter.

    HuggingFace dataset cards use YAML frontmatter at the top of README.md
    with fields like ``license``, ``task_categories``, ``language``, etc.
    """
    try:
        content = path.read_text(encoding="utf-8")
    except OSError as exc:
        logger.warning("Failed to read %s: %s", path, exc)
        return None

    match = _FRONTMATTER_RE.match(content)
    if not match:
        info = DatasetInfo(
            name=path.parent.name,
            source_file=str(path),
        )
        info.security_flags.append(
            {
                "severity": "LOW",
                "type": "NO_DATASET_CARD",
                "description": (
                    f"README.md in '{path.parent.name}' has no YAML frontmatter. Dataset cards improve transparency and compliance."
                ),
            }
        )
        return info

    frontmatter = match.group(1)
    kv = _parse_simple_yaml(frontmatter)

    info = DatasetInfo(
        name=kv.get("dataset_name", path.parent.name),
        license=kv.get("license", ""),
        source_file=str(path),
        version=kv.get("version", ""),
    )

    # Extract list fields
    info.task_categories = (
        _parse_yaml_list(frontmatter, "task_categories") or [kv.get("task_categories", "")] if kv.get("task_categories") else []
    )
    info.languages = _parse_yaml_list(frontmatter, "language") or [kv.get("language", "")] if kv.get("language") else []

    # Size category
    size_cat = kv.get("size_categories", "")
    if size_cat:
        info.description = f"Size category: {size_cat}"

    if not info.license:
        info.security_flags.append(
            {
                "severity": "MEDIUM",
                "type": "UNLICENSED_DATASET",
                "description": (
                    f"Dataset '{info.name}' has no license in YAML frontmatter. Training on unlicensed data may violate IP rights."
                ),
            }
        )

    return info


def parse_dvc_file(path: Path) -> DatasetInfo | None:
    """Parse a DVC ``.dvc`` tracking file for data provenance.

    DVC files are simple YAML with ``outs:`` containing ``md5``, ``size``, ``path``.
    """
    try:
        content = path.read_text(encoding="utf-8")
    except OSError as exc:
        logger.warning("Failed to read %s: %s", path, exc)
        return None

    # Extract the tracked file name from the .dvc filename (data.csv.dvc → data.csv)
    tracked_name = path.stem if path.suffix == ".dvc" else path.name

    info = DatasetInfo(
        name=tracked_name,
        source_file=str(path),
    )

    # Parse outs section for md5 and size
    md5_match = re.search(r"md5:\s+([a-f0-9]{32})", content)
    sha256_match = re.search(r"sha256:\s+([a-f0-9]{64})", content)
    size_match = re.search(r"size:\s+(\d+)", content)

    if md5_match:
        info.dvc_md5 = md5_match.group(1)
    elif sha256_match:
        info.dvc_md5 = sha256_match.group(1)

    if size_match:
        info.size_bytes = int(size_match.group(1))

    # Check for DVC remote config
    dvc_config = path.parent / ".dvc" / "config"
    if dvc_config.exists():
        try:
            config_text = dvc_config.read_text(encoding="utf-8")
            remote_match = re.search(r"url\s*=\s*(.+)", config_text)
            if remote_match:
                info.dvc_remote = remote_match.group(1).strip()
                info.security_flags.append(
                    {
                        "severity": "INFO",
                        "type": "REMOTE_DATA_SOURCE",
                        "description": f"DVC tracks data from remote: {info.dvc_remote[:100]}",
                    }
                )
        except OSError:
            pass

    if not info.dvc_md5:
        info.security_flags.append(
            {
                "severity": "LOW",
                "type": "UNVERSIONED_DATA",
                "description": f"DVC file '{tracked_name}' has no md5/sha256 hash. Data provenance cannot be verified.",
            }
        )

    return info


# ─── Discovery + batch scan ────────────────────────────────────────────────


def discover_dataset_files(directory: Path) -> list[Path]:
    """Find dataset card files in a directory tree."""
    results: list[Path] = []

    for p in directory.rglob("*"):
        if any(skip in p.parts for skip in _SKIP_DIRS):
            continue

        if p.name == "dataset_info.json":
            results.append(p)
        elif p.name == "dataset_infos.json":
            results.append(p)
        elif p.suffix == ".dvc" and p.is_file():
            results.append(p)
        elif p.name == "README.md" and p.is_file():
            # Only include READMEs that look like dataset cards (in data-like dirs)
            parent_name = p.parent.name.lower()
            # Heuristic: if sibling dataset_info.json exists, or parent looks like a dataset dir
            if (p.parent / "dataset_info.json").exists() or any(kw in parent_name for kw in ("data", "dataset", "corpus", "benchmark")):
                results.append(p)

    return sorted(results)


def scan_datasets(paths: list[Path]) -> DatasetScanResult:
    """Parse a list of dataset card files and return aggregated results."""
    result = DatasetScanResult()
    seen_names: set[str] = set()

    for path in paths:
        info: DatasetInfo | None = None

        if path.name in ("dataset_info.json", "dataset_infos.json"):
            info = parse_dataset_info_json(path)
        elif path.suffix == ".dvc":
            info = parse_dvc_file(path)
        elif path.name == "README.md":
            info = parse_dataset_readme(path)

        if info and info.name:
            # Deduplicate by name (prefer the one with more data)
            if info.name not in seen_names:
                tag_dataset(info)
                seen_names.add(info.name)
                result.datasets.append(info)
                result.source_files.append(str(path))

    return result


def scan_dataset_directory(directory: str | Path) -> DatasetScanResult:
    """Discover and scan all dataset card files in a directory."""
    d = Path(directory)
    if not d.is_dir():
        return DatasetScanResult(warnings=[f"Not a directory: {d}"])

    paths = discover_dataset_files(d)
    if not paths:
        return DatasetScanResult(warnings=[f"No dataset card files found in {d}"])

    return scan_datasets(paths)
