from __future__ import annotations

from collections.abc import Mapping, Sequence
import logging
from pathlib import Path
import re

from .json_files import (
    _infer_tuid_from_path,
    load_json_file_records,
    load_json_payload,
)
from ..utils.record_tools import (
    DEFAULT_KEEP_KEYS,
    EXTRA_FIELDS_KEY,
    EXTRA_TAGS_KEY,
    hybrid_qoi_records_from_payloads,
    split_record,
)
from ..write_data.write_batch import write_batch

logger = logging.getLogger(__name__)

AUTO_SNAPSHOT_FILE_PATTERNS = (
    "snapshot*.json",
    "snapshot*.json.bz2",
    "snapshot*.json.gz",
    "snapshot*.json.xz",
)
AUTO_QOI_FILE_PATTERNS = ("quantities_of_interest*.json",)
AUTO_EXCLUDED_FOLDER_NAMES = frozenset({"analysis_basicanalysis", ".ipynb_checkpoints"})
ELEMENT_TOKEN_PATTERN = re.compile(r"\bq\d+\b")


def ingest_data(
    folder: str | Path,
    *,
    continue_on_error: bool = True,
) -> list[dict[str, object]]:
    """Recursively load raw records from files in a folder.

    Parameters
    ----------
    folder : str | Path
        Folder to scan recursively.
    continue_on_error : bool, default=True
        If ``True``, skip unreadable files and continue.

    Returns
    -------
    list[dict[str, object]]
        Raw loaded records.

    Raises
    ------
    FileNotFoundError
        If folder does not exist.
    """
    root_path = Path(folder).expanduser()
    if not root_path.exists():
        raise FileNotFoundError(f"Folder not found: {root_path}")

    logger.info("ingest_data: scanning '%s'.", root_path)
    records = _ingest_auto_hybrid(root_path, continue_on_error=continue_on_error)
    logger.info("ingest_data: loaded %d raw records.", len(records))
    return records


def debug_ingest_matches(folder: str | Path) -> list[dict[str, str | None]]:
    """Print and return QoI-to-snapshot matching plan for recursive ingest.

    Parameters
    ----------
    folder : str | Path
        Folder scanned recursively.

    Returns
    -------
    list[dict[str, str | None]]
        List with ``qoi_path`` and ``snapshot_path`` (or ``None`` if unmatched).
    """
    root_path = Path(folder).expanduser()
    if not root_path.exists():
        raise FileNotFoundError(f"Folder not found: {root_path}")

    (
        snapshot_paths,
        qoi_paths,
        qoi_grouped_by_snapshot_path,
        _,
    ) = _collect_auto_hybrid_matches(root_path)

    root_resolved = root_path.resolve()
    snapshot_by_qoi: dict[Path, Path] = {}
    for snapshot_path, related_qoi_paths in qoi_grouped_by_snapshot_path.items():
        for qoi_path in related_qoi_paths:
            snapshot_by_qoi[qoi_path.resolve()] = snapshot_path.resolve()

    plan: list[dict[str, str | None]] = []
    if not qoi_paths:
        print("No quantities_of_interest*.json files found.")
        return plan

    for qoi_path in qoi_paths:
        qoi_resolved = qoi_path.resolve()
        matched_snapshot = snapshot_by_qoi.get(qoi_resolved)
        qoi_label = _display_path_for_debug(qoi_resolved, root_resolved)
        snapshot_label = (
            _display_path_for_debug(matched_snapshot, root_resolved)
            if matched_snapshot is not None
            else None
        )
        if snapshot_label is None:
            print(f"{qoi_label} -> (no snapshot, qoi-only)")
        else:
            print(f"{qoi_label} -> {snapshot_label}")
        plan.append(
            {
                "qoi_path": qoi_label,
                "snapshot_path": snapshot_label,
            }
        )

    snapshots_without_qoi = [
        path
        for path in snapshot_paths
        if path.resolve() not in set(snapshot_by_qoi.values())
    ]
    for snapshot_path in snapshots_without_qoi:
        snapshot_label = _display_path_for_debug(snapshot_path.resolve(), root_resolved)
        print(f"(snapshot-only) {snapshot_label}")

    return plan


def process_data(
    raw_data: Sequence[Mapping[str, object]],
    *,
    keep_keys: Sequence[str] | None = None,
    continue_on_error: bool = False,
) -> list[dict[str, object]]:
    """Normalize and validate raw records before writing.

    Parameters
    ----------
    raw_data : Sequence[Mapping[str, object]]
        Raw records from ``ingest_data``.
    keep_keys : Sequence[str] | None, default=None
        Keys to keep from each record. If ``None``, standard keys are used.
    continue_on_error : bool, default=False
        If ``True``, skip invalid records and continue.

    Returns
    -------
    list[dict[str, object]]
        Processed records ready for ``write_batch``.

    Raises
    ------
    ValueError
        If a record is invalid and ``continue_on_error`` is ``False``.
    """
    keep = set(keep_keys) if keep_keys is not None else set(DEFAULT_KEEP_KEYS)
    keep.update({"qoi", EXTRA_TAGS_KEY, EXTRA_FIELDS_KEY})

    processed: list[dict[str, object]] = []
    dropped = 0
    for raw_record in raw_data:
        try:
            filtered = {key: value for key, value in raw_record.items() if key in keep}
            record, extra_tags, extra_fields = split_record(filtered)
            _validate_processed_record(record, extra_fields)

            normalized: dict[str, object] = dict(record)
            if extra_tags:
                normalized[EXTRA_TAGS_KEY] = dict(extra_tags)
            if extra_fields:
                normalized[EXTRA_FIELDS_KEY] = dict(extra_fields)
            processed.append(normalized)
        except Exception as exc:
            dropped += 1
            logger.warning("Dropped record: %s", exc)
            if not continue_on_error:
                raise

    if dropped:
        logger.warning(
            "process_data: %d/%d records ready, %d dropped.",
            len(processed),
            len(raw_data),
            dropped,
        )
    else:
        logger.info(
            "process_data: %d/%d records ready.",
            len(processed),
            len(raw_data),
        )
    return processed


def ingest_and_write(
    folder: str | Path,
    *,
    keep_keys: Sequence[str] | None = None,
    settings_path: str | Path | None = None,
    continue_on_error: bool = True,
) -> dict[str, object]:
    """Ingest, process, and write records in one call.

    Parameters
    ----------
    folder : str | Path
        Folder to scan recursively.
    keep_keys : Sequence[str] | None, default=None
        Keys to keep before writing.
    settings_path : str | Path | None, default=None
        Optional settings TOML path.
    continue_on_error : bool, default=True
        If ``True``, continue after per-file, per-record, and write failures.

    Returns
    -------
    dict[str, object]
        Write summary plus ``raw_records_total`` and ``processed_records_total``.
    """
    logger.info("ingest_and_write: starting from '%s'.", folder)
    raw_data = ingest_data(
        folder,
        continue_on_error=continue_on_error,
    )
    processed_data = process_data(
        raw_data,
        keep_keys=keep_keys,
        continue_on_error=continue_on_error,
    )
    logger.info(
        "ingest_and_write: %d raw → %d processed records; writing now.",
        len(raw_data),
        len(processed_data),
    )
    write_summary = write_batch(
        processed_data,
        settings_path=settings_path,
        continue_on_error=continue_on_error,
    )

    summary = {
        "raw_records_total": len(raw_data),
        "processed_records_total": len(processed_data),
        **write_summary,
    }
    if write_summary.get("records_failed"):
        logger.warning(
            "ingest_and_write: complete — %d written, %d failed.",
            write_summary["records_written"],
            write_summary["records_failed"],
        )
    else:
        logger.info(
            "ingest_and_write: complete — %d records written.",
            write_summary["records_written"],
        )
    return summary


def _ingest_auto_hybrid(
    root_path: Path,
    *,
    continue_on_error: bool,
) -> list[dict[str, object]]:
    """Ingest Quantify-style data using QoI-first values plus snapshot metadata."""
    (
        snapshot_paths,
        _,
        qoi_grouped_by_snapshot_path,
        qoi_without_snapshot,
    ) = _collect_auto_hybrid_matches(root_path)

    records: list[dict[str, object]] = []
    for snapshot_path in snapshot_paths:
        related_qoi_paths = sorted(qoi_grouped_by_snapshot_path.get(snapshot_path, []))
        try:
            if related_qoi_paths:
                records.extend(
                    _load_hybrid_snapshot_qoi_records(
                        snapshot_path=snapshot_path,
                        qoi_paths=related_qoi_paths,
                    )
                )
            else:
                records.extend(load_json_file_records(snapshot_path))
        except Exception:
            logger.exception(
                "Failed parsing auto-ingest run with snapshot '%s'.",
                snapshot_path,
            )
            if not continue_on_error:
                raise

    for qoi_path in qoi_without_snapshot:
        try:
            records.extend(load_json_file_records(qoi_path))
        except Exception:
            logger.exception("Failed parsing file '%s'.", qoi_path)
            if not continue_on_error:
                raise

    return records


def _collect_auto_hybrid_matches(
    root_path: Path,
) -> tuple[list[Path], list[Path], dict[Path, list[Path]], list[Path]]:
    """Collect snapshot/QoI discovery and matching output for auto-ingest."""
    snapshot_paths = sorted(
        {
            path
            for pattern in AUTO_SNAPSHOT_FILE_PATTERNS
            for path in root_path.rglob(pattern)
            if path.is_file() and not _is_excluded_auto_ingest_path(path)
        }
    )
    qoi_paths = sorted(
        {
            path
            for pattern in AUTO_QOI_FILE_PATTERNS
            for path in root_path.rglob(pattern)
            if path.is_file() and not _is_excluded_auto_ingest_path(path)
        }
    )

    resolved_root = root_path.resolve()
    qoi_grouped_by_snapshot_path: dict[Path, list[Path]] = {}
    qoi_without_snapshot: list[Path] = []

    for qoi_path in qoi_paths:
        snapshot_path = _best_snapshot_match(
            qoi_path,
            snapshot_paths=snapshot_paths,
            root_path=resolved_root,
        )
        if snapshot_path is None:
            qoi_without_snapshot.append(qoi_path)
            continue
        qoi_grouped_by_snapshot_path.setdefault(snapshot_path, []).append(qoi_path)

    return snapshot_paths, qoi_paths, qoi_grouped_by_snapshot_path, qoi_without_snapshot


def _is_excluded_auto_ingest_path(path: Path) -> bool:
    """Return ``True`` when path is under excluded analysis folders."""
    return any(part.casefold() in AUTO_EXCLUDED_FOLDER_NAMES for part in path.parts)


def _best_snapshot_match(
    qoi_path: Path,
    *,
    snapshot_paths: Sequence[Path],
    root_path: Path,
) -> Path | None:
    """Select the best snapshot candidate for one QoI file."""
    qoi_parent = qoi_path.resolve().parent
    qoi_tuid = _infer_tuid_from_path(qoi_parent)

    ranked_candidates: list[tuple[tuple[int, int, int, str], Path]] = []
    for snapshot_path in snapshot_paths:
        snapshot_parent = snapshot_path.resolve().parent
        snapshot_tuid = _infer_tuid_from_path(snapshot_parent)
        if (
            qoi_tuid is not None
            and snapshot_tuid is not None
            and qoi_tuid != snapshot_tuid
        ):
            continue

        common_ancestor = _nearest_common_ancestor(
            qoi_parent,
            snapshot_parent,
            root_path=root_path,
        )
        if common_ancestor is None:
            continue

        qoi_distance = len(qoi_parent.relative_to(common_ancestor).parts)
        snapshot_distance = len(snapshot_parent.relative_to(common_ancestor).parts)
        is_direct_ancestor = 0 if snapshot_parent in qoi_path.resolve().parents else 1
        score = (
            is_direct_ancestor,
            qoi_distance + snapshot_distance,
            snapshot_distance,
            str(snapshot_path),
        )
        ranked_candidates.append((score, snapshot_path))

    if not ranked_candidates:
        return None
    ranked_candidates.sort(key=lambda item: item[0])
    return ranked_candidates[0][1]


def _nearest_common_ancestor(
    path_a: Path, path_b: Path, *, root_path: Path
) -> Path | None:
    """Return the nearest common ancestor under ``root_path``."""
    ancestors_b = {path_b, *path_b.parents}
    for candidate in (path_a, *path_a.parents):
        if candidate not in ancestors_b:
            continue
        if candidate == root_path or root_path in candidate.parents:
            return candidate
    return None


def _display_path_for_debug(path: Path, root_path: Path) -> str:
    """Return a stable debug label for a discovered file path."""
    try:
        return str(path.relative_to(root_path))
    except ValueError:
        return str(path)


def _load_hybrid_snapshot_qoi_records(
    *,
    snapshot_path: Path,
    qoi_paths: Sequence[Path],
) -> list[dict[str, object]]:
    """Merge QoI-map values with snapshot metadata for one run directory."""
    snapshot_payload = load_json_payload(snapshot_path)

    qoi_payloads: list[dict[str, object]] = []
    for qoi_path in qoi_paths:
        payload = load_json_payload(qoi_path)
        if isinstance(payload, Mapping):
            qoi_payloads.append({str(key): value for key, value in payload.items()})

    dataset_tuid = _infer_tuid_from_path(snapshot_path.parent)
    preferred_elements = _infer_element_names_from_path(snapshot_path.parent)
    records = hybrid_qoi_records_from_payloads(
        snapshot_payload={str(key): value for key, value in snapshot_payload.items()}
        if isinstance(snapshot_payload, Mapping)
        else {},
        qoi_payloads=qoi_payloads,
        dataset_tuid=dataset_tuid,
        preferred_elements=preferred_elements,
    )
    if records:
        return records
    return load_json_file_records(snapshot_path)


def _infer_element_names_from_path(path: Path) -> list[str]:
    """Infer qubit element names from folder labels such as ``... q2``."""
    names: set[str] = set()
    for part in path.parts:
        for match in ELEMENT_TOKEN_PATTERN.findall(part):
            names.add(match)
    return sorted(names)


def _validate_processed_record(
    record: Mapping[str, object],
    extra_fields: Mapping[str, object] | None,
) -> None:
    """Validate required pieces of a processed record."""
    qoi = record.get("qoi")
    if not isinstance(qoi, str) or not qoi.strip():
        raise ValueError("Each processed record must include a non-empty 'qoi'.")

    nominal_value = record.get("nominal_value")
    if isinstance(nominal_value, bool) or not isinstance(nominal_value, (int, float)):
        raise ValueError("Each processed record must include numeric 'nominal_value'.")

    uncertainty = record.get("uncertainty")
    if uncertainty is not None and (
        isinstance(uncertainty, bool) or not isinstance(uncertainty, (int, float))
    ):
        raise ValueError("'uncertainty' must be numeric when provided.")

    tuid = record.get("tuid")
    if tuid is not None and (not isinstance(tuid, str) or not tuid.strip()):
        raise ValueError("'tuid' must be a non-empty string when provided.")
