from __future__ import annotations

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

from .json_files import load_json_file_records
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
from ..write_data.writer import QOI_FIELD_KEYS

logger = logging.getLogger(__name__)

AUTO_JSON_FILE_PATTERNS = (
    "snapshot*.json",
    "quantities_of_interest*.json",
)
TUID_DIR_PREFIX_PATTERN = re.compile(
    r"^(?P<tuid>\d{8}-\d{6}-\d{3}-[A-Za-z0-9]{6})(?:-|$)"
)
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}")

    return _ingest_auto_hybrid(root_path, continue_on_error=continue_on_error)


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]] = []
    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:
            logger.exception("Failed processing one raw record.")
            if not continue_on_error:
                raise

    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, int]:
    """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, int]
        Write summary plus ``raw_records_total`` and ``processed_records_total``.
    """
    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,
    )
    write_summary = write_batch(
        processed_data,
        settings_path=settings_path,
        continue_on_error=continue_on_error,
    )

    return {
        "raw_records_total": len(raw_data),
        "processed_records_total": len(processed_data),
        **write_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 = sorted(
        path for path in root_path.rglob(AUTO_JSON_FILE_PATTERNS[0]) if path.is_file()
    )
    qoi_paths = sorted(
        path for path in root_path.rglob(AUTO_JSON_FILE_PATTERNS[1]) if path.is_file()
    )

    snapshot_by_dir = {path.parent.resolve(): path for path in snapshot_paths}
    qoi_grouped_by_snapshot_dir: dict[Path, list[Path]] = {}
    qoi_without_snapshot: list[Path] = []

    for qoi_path in qoi_paths:
        snapshot_dir = _nearest_snapshot_ancestor(
            qoi_path, snapshot_dirs=set(snapshot_by_dir)
        )
        if snapshot_dir is None:
            qoi_without_snapshot.append(qoi_path)
            continue
        qoi_grouped_by_snapshot_dir.setdefault(snapshot_dir, []).append(qoi_path)

    records: list[dict[str, object]] = []
    for snapshot_dir, snapshot_path in snapshot_by_dir.items():
        related_qoi_paths = sorted(qoi_grouped_by_snapshot_dir.get(snapshot_dir, []))
        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 _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 = json.loads(snapshot_path.read_text(encoding="utf-8"))

    qoi_payloads: list[dict[str, object]] = []
    for qoi_path in qoi_paths:
        payload = json.loads(qoi_path.read_text(encoding="utf-8"))
        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 _nearest_snapshot_ancestor(path: Path, *, snapshot_dirs: set[Path]) -> Path | None:
    """Return nearest ancestor containing a snapshot file."""
    for parent in path.resolve().parents:
        if parent in snapshot_dirs:
            return parent
    return None


def _infer_tuid_from_path(path: Path) -> str | None:
    """Infer Quantify TUID from a dataset directory name in the path."""
    for part in path.parts:
        match = TUID_DIR_PREFIX_PATTERN.match(part)
        if match is None:
            continue
        tuid = match.group("tuid").strip()
        if tuid:
            return tuid
    return None


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'.")

    has_builtin_field = any(record.get(key) is not None for key in QOI_FIELD_KEYS)
    has_extra_field = bool(extra_fields)
    if not (has_builtin_field or has_extra_field):
        raise ValueError(
            "Each processed record must include at least one field value "
            "or provide 'extra_fields'."
        )
