"""ELECTION-01 integrity gates (contract §5). No model fitting. MODEL_CHANGE=NO.

Executable negative-control primitives used by the harness. These are
synthetic/stub checks that prove interface separation and rejection rules —
not empirical election forecasts.

v0.2 responds to Codex/Astra adversarial review f4c7748b (source sha
9b8d33b6… preserved as integrity_gates_v0_failed_9b8d33b6.py):
1. Timezone-aware instant comparison (do not strip tz/time via date-only).
2. Sample collapse only on verified shared sample/tracker/duplicate links —
   coincident fieldwork alone does NOT merge independent houses.
3. Explicit train-before-test chronology guard (separate from set-disjointness).
4. Stub prediction invariance remains stub-only; production estimator must be
   wired through the same tests when available (see fit_eval_interface).
"""
from __future__ import annotations

from dataclasses import dataclass, field
from datetime import date, datetime, time, timezone
from typing import Any, Sequence
from zoneinfo import ZoneInfo

PIT_TZ = ZoneInfo("America/New_York")


def _parse_datetime(value: date | datetime | str) -> datetime | date:
    """Parse to datetime when timezone/time present; else date."""
    if isinstance(value, datetime):
        return value
    if isinstance(value, date):
        return value
    s = str(value).strip()
    if "T" in s or (len(s) > 10 and ("+" in s[10:] or s.endswith("Z"))):
        return datetime.fromisoformat(s.replace("Z", "+00:00"))
    return date.fromisoformat(s[:10])


def to_aware_instant(
    value: date | datetime | str,
    *,
    role: str,
    date_only_bound: str = "eod",
) -> datetime:
    """Convert to timezone-aware UTC-comparable instant.

    Date-only values use conservative America/New_York bounds:
      - cutoff / availability upper bound: end of day (23:59:59.999999 ET)
      - release / availability lower bound (asof start): start of day (00:00:00 ET)
    Never silently drop timezone info from aware datetimes.
    """
    parsed = _parse_datetime(value)
    if isinstance(parsed, datetime):
        if parsed.tzinfo is None:
            # Naive datetime: interpret in PIT_TZ (America/New_York), not local/UTC guess.
            return parsed.replace(tzinfo=PIT_TZ)
        return parsed
    # date-only
    if date_only_bound == "sod" or role in {"release", "available_asof"}:
        return datetime.combine(parsed, time(0, 0, 0), tzinfo=PIT_TZ)
    # cutoff / eod
    return datetime.combine(parsed, time(23, 59, 59, 999999), tzinfo=PIT_TZ)


def _as_date(value: date | datetime | str) -> date:
    """Fieldwork calendar date helper (date-only semantics intentionally)."""
    if isinstance(value, datetime):
        if value.tzinfo is not None:
            return value.astimezone(PIT_TZ).date()
        return value.date()
    if isinstance(value, date):
        return value
    s = str(value).strip()
    if "T" in s:
        return datetime.fromisoformat(s.replace("Z", "+00:00")).astimezone(PIT_TZ).date()
    return date.fromisoformat(s[:10])


@dataclass(frozen=True)
class RecordDecision:
    accepted: bool
    reason: str
    record_id: str


def reject_future_release(
    record_id: str,
    release_date: date | datetime | str,
    cutoff: date | datetime | str,
) -> RecordDecision:
    """§5.1 Future-release records are rejected when release_instant > cutoff_instant.

    Compares timezone-aware instants. Date-only cutoff = ET end-of-day;
    date-only release = ET start-of-day (conservative: available at beginning of that day).
    """
    rel = to_aware_instant(release_date, role="release", date_only_bound="sod")
    cut = to_aware_instant(cutoff, role="cutoff", date_only_bound="eod")
    if rel > cut:
        return RecordDecision(False, "future_release", record_id)
    return RecordDecision(True, "ok", record_id)


def reject_unavailable_revision(
    record_id: str,
    available_asof: date | datetime | str | None,
    cutoff: date | datetime | str,
    *,
    known_unavailable: bool = False,
) -> RecordDecision:
    """§5.2 Unavailable revisions are rejected.

    Rejects when available_asof is missing, marked unavailable, or after cutoff.
    Instant comparison — does not strip timezone/time.
    """
    if known_unavailable or available_asof is None:
        return RecordDecision(False, "unavailable_revision", record_id)
    avail = to_aware_instant(available_asof, role="available_asof", date_only_bound="sod")
    cut = to_aware_instant(cutoff, role="cutoff", date_only_bound="eod")
    if avail > cut:
        return RecordDecision(False, "unavailable_revision", record_id)
    return RecordDecision(True, "ok", record_id)


def assert_no_train_test_cross(
    train_elections: Sequence[str | int],
    test_elections: Sequence[str | int],
) -> dict[str, Any]:
    """§5.4 No election crosses training/test boundaries (set disjointness only).

    This does NOT enforce chronology. Use assert_train_before_test for ordering.
    """
    train = {str(x) for x in train_elections}
    test = {str(x) for x in test_elections}
    overlap = sorted(train & test)
    return {
        "ok": len(overlap) == 0,
        "overlap": overlap,
        "train": sorted(train),
        "test": sorted(test),
        "role": "set_disjointness_only",
        "chronology_enforced": False,
    }


def assert_train_before_test(
    train_elections: Sequence[str | int],
    test_elections: Sequence[str | int],
) -> dict[str, Any]:
    """Explicit chronology gate: max(train) < min(test) on numeric year ids.

    Non-numeric ids are reported; chronology check requires all-numeric years.
    """
    def as_years(xs: Sequence[str | int]) -> tuple[list[int], list[str]]:
        years: list[int] = []
        bad: list[str] = []
        for x in xs:
            s = str(x)
            try:
                years.append(int(s))
            except ValueError:
                bad.append(s)
        return years, bad

    train_y, train_bad = as_years(train_elections)
    test_y, test_bad = as_years(test_elections)
    if train_bad or test_bad or not train_y or not test_y:
        return {
            "ok": False,
            "reason": "non_numeric_or_empty",
            "train_bad": train_bad,
            "test_bad": test_bad,
            "train": [str(x) for x in train_elections],
            "test": [str(x) for x in test_elections],
        }
    max_train = max(train_y)
    min_test = min(test_y)
    return {
        "ok": max_train < min_test,
        "max_train": max_train,
        "min_test": min_test,
        "train": sorted(train_y),
        "test": sorted(test_y),
        "role": "chronology_max_train_lt_min_test",
    }


@dataclass(frozen=True)
class PollSample:
    poll_id: str
    sample_id: str
    field_start: date | datetime | str
    field_end: date | datetime | str
    house: str = ""
    source: str = ""
    tracker_id: str | None = None
    duplicate_of: str | None = None  # scoped sample key this duplicates
    verified_shared_with: tuple[str, ...] = ()  # scoped sample keys

    def scoped_sample_key(self) -> str:
        src = self.source.strip() or "_default"
        return f"{src}::{self.sample_id}"


def _overlap_days(a0: date, a1: date, b0: date, b1: date) -> int:
    start = max(a0, b0)
    end = min(a1, b1)
    if end < start:
        return 0
    return (end - start).days + 1


def collapse_overlapping_samples(
    polls: Sequence[PollSample],
    *,
    min_overlap_days: int = 1,
) -> dict[str, Any]:
    """§5.5 Verified-shared samples are not counted as independent polls.

    Rules:
      - Each scoped sample_id (source::sample_id) is at most one independent unit.
      - Distinct sample keys merge ONLY with verified shared relationships:
        same non-empty tracker_id, duplicate_of link, or verified_shared_with.
      - Coincident / overlapping fieldwork alone is NOT evidence of respondent
        overlap across independent houses — those are flagged as
        uncertain_dependence, not merged.
    """
    by_sample: dict[str, list[PollSample]] = {}
    for p in polls:
        by_sample.setdefault(p.scoped_sample_key(), []).append(p)

    sample_reps: list[tuple[str, PollSample, list[PollSample]]] = []
    for sample_key, members in sorted(by_sample.items()):
        rep = sorted(members, key=lambda m: m.poll_id)[0]
        sample_reps.append((sample_key, rep, list(members)))

    n = len(sample_reps)
    parent = list(range(n))

    def find(i: int) -> int:
        while parent[i] != i:
            parent[i] = parent[parent[i]]
            i = parent[i]
        return i

    def union(i: int, j: int) -> None:
        ri, rj = find(i), find(j)
        if ri != rj:
            parent[rj] = ri

    key_to_idx = {sample_reps[i][0]: i for i in range(n)}

    # Build verified-share edges
    for i, (key, rep, members) in enumerate(sample_reps):
        # tracker_id: merge samples sharing the same non-empty tracker
        trackers = {m.tracker_id for m in members if m.tracker_id}
        for j, (key2, rep2, members2) in enumerate(sample_reps):
            if j <= i:
                continue
            trackers2 = {m.tracker_id for m in members2 if m.tracker_id}
            if trackers and trackers2 and trackers & trackers2:
                union(i, j)
        # duplicate_of / verified_shared_with
        for m in members:
            if m.duplicate_of and m.duplicate_of in key_to_idx:
                union(i, key_to_idx[m.duplicate_of])
            for other in m.verified_shared_with:
                if other in key_to_idx:
                    union(i, key_to_idx[other])

    spans = []
    for sample_key, rep, members in sample_reps:
        starts = [_as_date(m.field_start) for m in members]
        ends = [_as_date(m.field_end) for m in members]
        spans.append((min(starts), max(ends), rep.house))

    uncertain: list[dict[str, Any]] = []
    for i in range(n):
        for j in range(i + 1, n):
            if find(i) == find(j):
                continue
            a0, a1, house_a = spans[i]
            b0, b1, house_b = spans[j]
            ov = _overlap_days(a0, a1, b0, b1)
            if ov >= min_overlap_days:
                uncertain.append(
                    {
                        "sample_keys": [sample_reps[i][0], sample_reps[j][0]],
                        "houses": [house_a, house_b],
                        "overlap_days": ov,
                        "action": "flag_uncertain_dependence_not_merged",
                    }
                )

    clusters: dict[int, list[int]] = {}
    for i in range(n):
        clusters.setdefault(find(i), []).append(i)

    independent: list[str] = []
    groups: list[dict[str, Any]] = []
    for root, idxs in sorted(clusters.items()):
        all_members: list[PollSample] = []
        sample_ids = []
        for i in idxs:
            sid, rep, members = sample_reps[i]
            sample_ids.append(sid)
            all_members.extend(members)
        rep = sorted(all_members, key=lambda m: m.poll_id)[0]
        independent.append(rep.poll_id)
        groups.append(
            {
                "sample_keys": sorted(sample_ids),
                "representative_poll_id": rep.poll_id,
                "member_poll_ids": sorted(m.poll_id for m in all_members),
                "collapsed_count": len(all_members),
            }
        )

    raw_count = len(polls)
    indep_count = len(independent)
    return {
        "ok": True,
        "raw_poll_count": raw_count,
        "independent_count": indep_count,
        "independent_poll_ids": independent,
        "groups": groups,
        "uncertain_dependence": uncertain,
        "no_false_independence": indep_count <= raw_count,
        "merge_rule": "verified_shared_only",
    }


def identical_eligible_test_cycles(
    cycles_a: Sequence[str | int],
    cycles_b: Sequence[str | int],
) -> dict[str, Any]:
    """§5.6 Comparisons use identical eligible test cycles."""
    a = [str(x) for x in cycles_a]
    b = [str(x) for x in cycles_b]
    set_a, set_b = set(a), set(b)
    return {
        "ok": a == b,  # identical ordered eligible list
        "equal_as_sets": set_a == set_b,
        "cycles_a": a,
        "cycles_b": b,
        "only_in_a": sorted(set_a - set_b),
        "only_in_b": sorted(set_b - set_a),
    }


def utc_now_iso() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")