"""CLI utilities for inspecting prepared dataset partitions."""

from __future__ import annotations

import argparse
import csv
import json
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Mapping, Optional, Sequence, Set

from idxr.prepare_datasets_lib.manifest import (
    MANIFEST_FILENAME,
    load_manifest,
)

DEFAULT_PARTITIONS_DIR = Path("workdir/partitions")
DEFAULT_INSIGHTS_DIR = Path("dataset_insights")
DEFAULT_SUMMARY_OUTPUT = Path("dataset_summary.json")


@dataclass
class FieldAccumulator:
    """Tracks statistics for a single field."""

    unique_values: Set[str] = field(default_factory=set)
    observations: int = 0
    null_count: int = 0
    partitions_missing_column: Set[str] = field(default_factory=set)
    rows_missing_column: int = 0

    def record_value(self, raw_value: str) -> None:
        """Track a single observed value."""
        self.observations += 1
        value = raw_value.strip()
        if value == "":
            self.null_count += 1
        else:
            self.unique_values.add(value)

    def register_missing_partition(
        self,
        partition_name: str,
        missing_rows: int,
    ) -> None:
        """Mark that a partition omitted this column entirely."""
        self.partitions_missing_column.add(partition_name)
        self.rows_missing_column += missing_rows

    def as_insights(self) -> Dict[str, Any]:
        """Serialize detailed insights for JSON output."""
        non_null = self.observations - self.null_count
        return {
            "unique_values": sorted(self.unique_values),
            "unique_count": len(self.unique_values),
            "observations": self.observations,
            "non_null_count": non_null,
            "null_count": self.null_count,
            "rows_missing_column": self.rows_missing_column,
            "partitions_missing_column": sorted(self.partitions_missing_column),
        }

    def as_summary(self) -> Dict[str, Any]:
        """Serialize summary statistics for JSON output."""
        non_null = self.observations - self.null_count
        return {
            "unique_count": len(self.unique_values),
            "non_null_count": non_null,
            "null_count": self.null_count,
            "observations": self.observations,
            "missing_partition_count": len(self.partitions_missing_column),
            "rows_missing_column": self.rows_missing_column,
        }


@dataclass
class ModelAccumulator:
    """Aggregates statistics for a model."""

    fields: Dict[str, FieldAccumulator] = field(default_factory=dict)
    rows: int = 0

    def ensure_field(self, name: str) -> FieldAccumulator:
        """Ensure a field accumulator exists."""
        if name not in self.fields:
            self.fields[name] = FieldAccumulator()
        return self.fields[name]

    def as_insights(self) -> Dict[str, Any]:
        return {
            "rows_analyzed": self.rows,
            "fields": {
                column: stats.as_insights() for column, stats in self.fields.items()
            },
        }

    def as_summary(self) -> Dict[str, Any]:
        return {
            "rows_analyzed": self.rows,
            "field_count": len(self.fields),
            "fields": {
                column: stats.as_summary() for column, stats in self.fields.items()
            },
        }


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="idxr dataset_insights",
        description=(
            "Scan prepare_datasets partitions and capture unique-value insights for "
            "selected models/fields."
        ),
    )
    parser.add_argument(
        "--partitions-dir",
        type=Path,
        default=DEFAULT_PARTITIONS_DIR,
        help=(
            "Directory containing partition_* subdirectories and a manifest.json produced "
            "by prepare_datasets (default: %(default)s)."
        ),
    )
    parser.add_argument(
        "--fields-config",
        type=Path,
        default=None,
        help=(
            "Path to a JSON file mapping model names to the fields that should be "
            "inspected for unique values. If omitted, every column present in each "
            "model CSV will be analyzed."
        ),
    )
    parser.add_argument(
        "--insights-dir",
        type=Path,
        default=DEFAULT_INSIGHTS_DIR,
        help=(
            "Directory where per-model insight files will be written "
            "(default: %(default)s)."
        ),
    )
    parser.add_argument(
        "--summary-output",
        type=Path,
        default=DEFAULT_SUMMARY_OUTPUT,
        help="Where to write the summary statistics JSON (default: %(default)s).",
    )
    parser.add_argument(
        "--log-level",
        default="INFO",
        choices={"DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"},
        help="Logging verbosity (default: %(default)s).",
    )
    return parser


def _load_field_config(config_path: Optional[Path]) -> Optional[Dict[str, Set[str]]]:
    if config_path is None:
        return None
    if not config_path.exists():
        raise FileNotFoundError(f"Fields configuration not found at {config_path}")
    data = json.loads(config_path.read_text(encoding="utf-8"))
    if not isinstance(data, Mapping):
        raise ValueError("Fields configuration must be a JSON object.")

    mapping: Dict[str, Set[str]] = {}
    for model_name, raw_fields in data.items():
        if not isinstance(model_name, str):
            continue
        if not isinstance(raw_fields, list):
            logging.warning(
                "Skipping %s entry because the field list is not an array.",
                model_name,
            )
            continue
        fields = {
            str(column) for column in raw_fields if isinstance(column, str) and column
        }
        if fields:
            mapping[model_name] = fields
    return mapping


def _iter_partitions(manifest: Mapping[str, Any]) -> Sequence[Mapping[str, Any]]:
    partitions = manifest.get("partitions", [])
    if not isinstance(partitions, list):
        raise ValueError("Manifest 'partitions' entry must be a list.")
    return [
        entry
        for entry in partitions
        if isinstance(entry, Mapping) and not entry.get("stale", False)
    ]


def _read_csv(
    *,
    csv_path: Path,
    selected_fields: Optional[Set[str]],
    accumulator: ModelAccumulator,
    partition_name: str,
) -> None:
    if selected_fields is not None and not selected_fields:
        return
    if not csv_path.exists():
        logging.warning("Skipping missing CSV for %s: %s", partition_name, csv_path)
        return

    try:
        with csv_path.open("r", encoding="utf-8", newline="") as handle:
            reader = csv.reader(handle)
            headers = next(reader, None)
            if not headers:
                logging.info("Skipping empty CSV %s", csv_path)
                return
            header_to_index = {name: idx for idx, name in enumerate(headers)}
            header_names = list(header_to_index.keys())
            header_set = set(header_names)

            if selected_fields is None:
                for column in header_names:
                    accumulator.ensure_field(column)
                target_fields = header_names
                expected_fields = set(accumulator.fields.keys())
            else:
                for column in selected_fields:
                    accumulator.ensure_field(column)
                target_fields = [
                    column for column in selected_fields if column in header_to_index
                ]
                expected_fields = set(selected_fields)

            if not target_fields and not expected_fields:
                return

            rows_read = 0
            for row in reader:
                rows_read += 1
                for column in target_fields:
                    idx = header_to_index.get(column)
                    if idx is None:
                        continue
                    value = row[idx] if idx < len(row) else ""
                    accumulator.fields[column].record_value(value)

            accumulator.rows += rows_read
            if rows_read == 0:
                return

            missing_fields = expected_fields - header_set
            if missing_fields:
                for column in missing_fields:
                    accumulator.fields[column].register_missing_partition(
                        partition_name, rows_read
                    )
    except OSError as exc:
        logging.error("Failed to read %s: %s", csv_path, exc)


def collect_insights(
    *,
    partitions_dir: Path,
    field_map: Optional[Mapping[str, Set[str]]],
) -> Dict[str, ModelAccumulator]:
    manifest_path = partitions_dir / MANIFEST_FILENAME
    manifest = load_manifest(manifest_path)
    model_stats: Dict[str, ModelAccumulator] = {}

    partitions = _iter_partitions(manifest)
    logging.info("Analyzing %s partition(s) in %s", len(partitions), partitions_dir)
    for partition in partitions:
        name = partition.get("name") or "unknown-partition"
        partition_models = partition.get("models", {})
        if not isinstance(partition_models, Mapping):
            logging.debug(
                "Skipping partition %s because 'models' is not a mapping", name
            )
            continue

        if field_map:
            target_models = field_map.keys()
        else:
            target_models = partition_models.keys()

        for model_name in target_models:
            configured_fields = field_map.get(model_name) if field_map else None
            if configured_fields is not None and not configured_fields:
                continue
            model_info = partition_models.get(model_name)
            if not isinstance(model_info, Mapping):
                continue
            path_value = model_info.get("path")
            if not isinstance(path_value, str):
                logging.debug(
                    "Skipping partition %s model %s because path is missing.",
                    name,
                    model_name,
                )
                continue
            accumulator = model_stats.get(model_name)
            if accumulator is None:
                accumulator = ModelAccumulator(
                    fields=(
                        {column: FieldAccumulator() for column in configured_fields}
                        if configured_fields
                        else {}
                    )
                )
                model_stats[model_name] = accumulator
            _read_csv(
                csv_path=Path(path_value),
                selected_fields=configured_fields,
                accumulator=accumulator,
                partition_name=name,
            )
    return model_stats


def write_outputs(
    *,
    stats: Mapping[str, ModelAccumulator],
    insights_dir: Path,
    summary_output: Path,
) -> None:
    insights_dir.mkdir(parents=True, exist_ok=True)
    summary_payload: Dict[str, Any] = {
        "models": {
            model: accumulator.as_summary() for model, accumulator in stats.items()
        }
    }
    insights_files: Dict[str, str] = {}
    for model, accumulator in stats.items():
        model_path = insights_dir / f"{model}.json"
        model_path.write_text(
            json.dumps(accumulator.as_insights(), indent=2), encoding="utf-8"
        )
        insights_files[model] = str(model_path.resolve())

    summary_payload["insights_files"] = insights_files

    summary_output.parent.mkdir(parents=True, exist_ok=True)
    summary_output.write_text(json.dumps(summary_payload, indent=2), encoding="utf-8")


def main(argv: Sequence[str] | None = None) -> int:
    parser = _build_parser()
    args = parser.parse_args(list(argv) if argv is not None else None)
    logging.basicConfig(level=getattr(logging, args.log_level))

    try:
        field_map = _load_field_config(args.fields_config)
    except (OSError, ValueError) as exc:
        parser.error(str(exc))

    stats = collect_insights(
        partitions_dir=args.partitions_dir,
        field_map=field_map,
    )
    if not stats:
        logging.warning("No dataset insights were collected.")
        return 1

    write_outputs(
        stats=stats,
        insights_dir=args.insights_dir,
        summary_output=args.summary_output,
    )
    logging.info(
        "Wrote per-model insights to %s and summary to %s",
        args.insights_dir,
        args.summary_output,
    )
    return 0


if __name__ == "__main__":  # pragma: no cover
    raise SystemExit(main())
