from __future__ import annotations

import hashlib
from collections import defaultdict
from typing import Any, Iterable, Mapping, Sequence

import numpy as np

from .features import FEATURE_NAMES, feature_vector


REQUIRED_RECORD_FIELDS = (
    "sample_id",
    "field_id",
    "flight_id",
    "capture_date",
    "data_provenance",
    "label",
    "label_reference",
    "features",
)
DEFAULT_GROUP_FIELDS = ("field_id", "flight_id", "capture_date")
SPLIT_NAMES = ("train", "validation", "test")


def validate_record(record: Mapping[str, Any]) -> None:
    missing = [name for name in REQUIRED_RECORD_FIELDS if name not in record]
    if missing:
        raise ValueError(f"Dataset record is missing fields: {missing}")
    if not isinstance(record["sample_id"], str) or not record["sample_id"]:
        raise ValueError("sample_id must be a non-empty string")
    if int(record["label"]) not in {0, 1}:
        raise ValueError("label must be binary 0 or 1")
    if record["data_provenance"] not in {"synthetic", "field"}:
        raise ValueError("data_provenance must be synthetic or field")
    if record["label_reference"] not in {
        "synthetic_latent_truth",
        "visible_symptom_review",
        "expert_consensus",
        "laboratory_confirmed",
    }:
        raise ValueError("label_reference is not in the published evidence hierarchy")
    feature_vector(record["features"])


def validate_records(records: Sequence[Mapping[str, Any]]) -> None:
    if not records:
        raise ValueError("Dataset must contain at least one record")
    sample_ids: set[str] = set()
    for record in records:
        validate_record(record)
        sample_id = str(record["sample_id"])
        if sample_id in sample_ids:
            raise ValueError(f"Duplicate sample_id: {sample_id}")
        sample_ids.add(sample_id)


def group_key(record: Mapping[str, Any], group_fields: Iterable[str] = DEFAULT_GROUP_FIELDS) -> str:
    fields = tuple(group_fields)
    if not fields:
        raise ValueError("At least one group field is required")
    values: list[str] = []
    for field in fields:
        value = record.get(field)
        if not isinstance(value, str) or not value:
            raise ValueError(f"Grouping field {field!r} is missing or empty")
        values.append(value)
    return "|".join(values)


def split_by_group(
    records: Sequence[Mapping[str, Any]],
    *,
    seed: int = 20260804,
    group_fields: Iterable[str] = DEFAULT_GROUP_FIELDS,
    validation_fraction: float = 0.20,
    test_fraction: float = 0.20,
) -> dict[str, str]:
    validate_records(records)
    if validation_fraction <= 0 or test_fraction <= 0 or validation_fraction + test_fraction >= 1:
        raise ValueError("Split fractions must be positive and leave a training partition")
    fields = tuple(group_fields)
    groups = sorted({group_key(record, fields) for record in records})
    if len(groups) < 3:
        raise ValueError("At least three independent groups are required")
    order = np.random.default_rng(seed).permutation(len(groups))
    shuffled = [groups[int(index)] for index in order]
    number_test = max(1, int(round(len(groups) * test_fraction)))
    number_validation = max(1, int(round(len(groups) * validation_fraction)))
    if number_test + number_validation >= len(groups):
        number_test = 1
        number_validation = 1
    test_groups = set(shuffled[:number_test])
    validation_groups = set(shuffled[number_test:number_test + number_validation])
    return {
        group: "test" if group in test_groups else "validation" if group in validation_groups else "train"
        for group in groups
    }


def assign_splits(
    records: Sequence[Mapping[str, Any]],
    assignments: Mapping[str, str],
    *,
    group_fields: Iterable[str] = DEFAULT_GROUP_FIELDS,
) -> list[dict[str, Any]]:
    fields = tuple(group_fields)
    result: list[dict[str, Any]] = []
    for original in records:
        key = group_key(original, fields)
        split = assignments.get(key)
        if split not in SPLIT_NAMES:
            raise ValueError(f"Group {key} has no valid split assignment")
        row = dict(original)
        row["group_key"] = key
        row["split"] = split
        result.append(row)
    assert_no_group_leakage(result)
    return result


def assert_no_group_leakage(records: Sequence[Mapping[str, Any]]) -> None:
    observed: dict[str, set[str]] = defaultdict(set)
    for record in records:
        key = record.get("group_key")
        split = record.get("split")
        if not isinstance(key, str) or split not in SPLIT_NAMES:
            raise ValueError("Split records must contain group_key and a valid split")
        observed[key].add(str(split))
    leaking = sorted(key for key, splits in observed.items() if len(splits) > 1)
    if leaking:
        raise ValueError(f"Group leakage detected: {leaking}")


def matrix_from_records(
    records: Sequence[Mapping[str, Any]],
    *,
    split: str | None = None,
) -> tuple[np.ndarray, np.ndarray, list[str]]:
    selected = [record for record in records if split is None or record.get("split") == split]
    if not selected:
        raise ValueError(f"No records available for split {split!r}")
    matrix = np.vstack([feature_vector(record["features"]) for record in selected])
    labels = np.asarray([int(record["label"]) for record in selected], dtype=np.int64)
    return matrix, labels, [str(record["sample_id"]) for record in selected]


def split_receipt(
    records: Sequence[Mapping[str, Any]],
    assignments: Mapping[str, str],
    *,
    group_fields: Iterable[str] = DEFAULT_GROUP_FIELDS,
    seed: int,
) -> dict[str, Any]:
    assigned = assign_splits(records, assignments, group_fields=group_fields)
    counts = {
        split: {
            "samples": sum(row["split"] == split for row in assigned),
            "groups": len({row["group_key"] for row in assigned if row["split"] == split}),
            "positive": sum(row["split"] == split and int(row["label"]) == 1 for row in assigned),
            "negative": sum(row["split"] == split and int(row["label"]) == 0 for row in assigned),
        }
        for split in SPLIT_NAMES
    }
    digest = hashlib.sha256()
    for row in sorted(assigned, key=lambda value: str(value["sample_id"])):
        digest.update(f"{row['sample_id']}\0{row['group_key']}\0{row['split']}\n".encode("utf-8"))
    return {
        "schemaVersion": 1,
        "seed": seed,
        "groupFields": list(group_fields),
        "leakageDetected": False,
        "counts": counts,
        "assignmentSha256": digest.hexdigest(),
        "interpretation": "Groups, not pixels, were split. Synthetic split performance is engineering verification, not field disease performance.",
    }


def assess_model_promotion(records: Sequence[Mapping[str, Any]]) -> dict[str, Any]:
    """Fail closed unless a held-out field has laboratory-confirmed truth.

    This gate is intentionally stricter than software-test success. Passing it
    only permits consideration for field validation; it never authorizes
    autonomous treatment or a production disease claim.
    """
    assert_no_group_leakage(records)
    failures: list[str] = []
    if any(record.get("data_provenance") != "field" for record in records):
        failures.append("non_field_or_synthetic_data_present")
    train_fields = {str(record["field_id"]) for record in records if record.get("split") == "train"}
    test_fields = {str(record["field_id"]) for record in records if record.get("split") == "test"}
    if not train_fields or not test_fields or train_fields & test_fields:
        failures.append("no_disjoint_held_out_field")
    train_dates = {str(record["capture_date"]) for record in records if record.get("split") == "train"}
    test_dates = {str(record["capture_date"]) for record in records if record.get("split") == "test"}
    if not train_dates or not test_dates or train_dates & test_dates:
        failures.append("no_disjoint_held_out_date")
    laboratory_test = [
        record for record in records
        if record.get("split") == "test" and record.get("label_reference") == "laboratory_confirmed"
    ]
    if {int(record["label"]) for record in laboratory_test} != {0, 1}:
        failures.append("held_out_test_lacks_both_laboratory_classes")
    if len(laboratory_test) < 60:
        failures.append("fewer_than_60_independent_laboratory_test_units")
    if any(
        not record.get("second_coder_id") or record.get("adjudication_status") not in {"not_required", "resolved"}
        for record in laboratory_test
    ):
        failures.append("held_out_test_lacks_second_coder_and_adjudication")
    status = "FIELD_VALIDATION_CANDIDATE" if not failures else "RESEARCH_ONLY"
    return {
        "schemaVersion": 1,
        "status": status,
        "failures": sorted(set(failures)),
        "laboratoryHeldOutUnits": len(laboratory_test),
        "trainFields": sorted(train_fields),
        "testFields": sorted(test_fields),
        "trainDates": sorted(train_dates),
        "testDates": sorted(test_dates),
        "boundary": "This gate does not establish diagnostic validity or authorize autonomous flight, pesticide use, or treatment.",
    }


def records_manifest(records: Sequence[Mapping[str, Any]]) -> dict[str, Any]:
    validate_records(records)
    provenance = sorted({str(record["data_provenance"]) for record in records})
    references = sorted({str(record["label_reference"]) for record in records})
    return {
        "schemaVersion": 1,
        "featureContract": list(FEATURE_NAMES),
        "sampleCount": len(records),
        "fieldCount": len({str(record["field_id"]) for record in records}),
        "groupCount": len({group_key(record) for record in records}),
        "positiveCount": sum(int(record["label"]) == 1 for record in records),
        "negativeCount": sum(int(record["label"]) == 0 for record in records),
        "dataProvenance": provenance,
        "labelReferences": references,
        "containsFieldPerformance": False,
        "warning": "Synthetic labels and metrics verify software only; they are not field plant-disease performance.",
    }
