#!/usr/bin/env python3
"""Select one bounded, source-backed module batch for high-throughput parity work.

The planner does not change dispositions or task status. It only materializes a
small deterministic work packet so one implementation turn can close several
records that share a native adapter family.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
from collections import Counter, defaultdict
from datetime import datetime, timezone
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from implementation_first import all_implementation_ready_ids, implementation_first_enabled  # noqa: E402


ROOT = Path(__file__).resolve().parents[1]
DOCS = ROOT / "docs"
P6 = DOCS / "p6"
STATE = DOCS / "current/task-state.json"
OUTPUT = DOCS / "current/module-batch-manifest.json"
SCHEDULER_STATUS = DOCS / "current/freecad-web-scheduler-status.json"
SCHEDULER_TIMING = DOCS / "current/freecad-web-scheduler-timing.json"
TASK_STATE = DOCS / "current/task-state.json"
CONTINUATION_STATE = DOCS / "current/continuation-state.json"
EVIDENCE_INDEX = DOCS / "current/module-batch-evidence-index.json"

PREFERRED_MODULES = {
    "Part": 0,
    "PartDesign": 1,
    "Sketcher": 2,
    "TechDraw": 3,
    "Draft": 4,
    "Material": 5,
    "Mesh": 6,
    "Import": 7,
    "BIM": 8,
    "Assembly": 9,
    "Fem": 10,
}
TIER_ORDER = {"tier-1-readonly": 0, "tier-2-pure": 1, "tier-3-stateful": 2, "tier-4-restricted": 3}
FAMILY_ORDER = {
    "pure-function": 0,
    "geometry-value": 1,
    "generic-object": 2,
    "property": 3,
    "document-graph": 4,
    "serialization": 5,
    "lifecycle-command": 6,
    "view-resource": 7,
}
UNRESOLVED = {"planned", "not_implemented"}
TERMINAL_RECORD_STATES = {
    "exposed", "restricted", "nonportable", "blocked_native_probe",
    "passed", "passed_boundary", "pass", "completed", "native_handler",
}
_PENDING_EVIDENCE_CACHE: set[str] | None = None


def adaptive_target_size(requested: int | None, tier: str, module: str, family: str) -> tuple[int, str]:
    """Adjust a family using recent complete-cycle throughput and failures."""
    base = requested if requested is not None else 24
    if requested is not None:
        return max(4, min(64, requested)), "explicit"
    reason = "default"
    try:
        timing = json.loads(SCHEDULER_TIMING.read_text(encoding="utf-8"))
        samples: list[tuple[float, int, bool]] = []
        for cycle in reversed(timing.get("cycles", [])):
            batch = cycle.get("batch", {}) if isinstance(cycle, dict) else {}
            if (batch.get("module"), batch.get("tier"), batch.get("family")) != (module, tier, family):
                continue
            duration = float(
                cycle.get(
                    "effectiveCycleSeconds",
                    float(cycle.get("turnSeconds", 0) or 0)
                    + float(cycle.get("preparationSeconds", 0) or 0)
                    + float(cycle.get("postPipelineSeconds", 0) or 0),
                ) or 0
            )
            activity = cycle.get("activity", {}) if isinstance(cycle.get("activity"), dict) else {}
            failed = bool(cycle.get("timedOut")) or int(activity.get("commandFailures", 0) or 0) > 0 or bool(cycle.get("commandBudgetExceeded"))
            samples.append((duration, max(4, int(batch.get("records", base) or base)), failed))
            if len(samples) == 3:
                break
        if samples:
            durations = [value[0] for value in samples if value[0] > 0]
            previous_size = samples[0][1]
            if any(value[2] for value in samples):
                return max(4, previous_size // 2), "same_family_recent_failure"
            if durations:
                average = sum(durations) / len(durations)
                maximum = max(durations)
                records_per_second = sum(value[1] for value in samples) / sum(durations)
                if maximum > 900:
                    return max(4, previous_size // 2), "same_family_effective_over_900s"
                if average > 420:
                    return max(4, previous_size - 4), "same_family_average_over_420s"
                if average < 90 and records_per_second >= 0.25:
                    return min(32, previous_size + 8), "same_family_three_cycle_fast_throughput"
                if average < 240:
                    return min(32, previous_size + 4), "same_family_three_cycle_under_240s"
    except (OSError, json.JSONDecodeError, TypeError, ValueError):
        pass
    try:
        status = json.loads(SCHEDULER_STATUS.read_text(encoding="utf-8"))
        previous = (status.get("moduleBatchModule"), status.get("moduleBatchTier"), status.get("moduleBatchFamily"))
        current = (module, tier, family)
        duration = float(status.get("lastTurnDurationSeconds", 0) or 0) + float(status.get("postTurnPipelineSeconds", 0) or 0)
        timed_out = bool(status.get("timedOut"))
        if previous == current and (timed_out or duration > 900):
            base = max(4, base // 2)
            reason = "previous_timeout_or_over_900s"
        elif previous == current and duration and duration < 240 and not timed_out:
            base = min(32, base + 4)
            reason = "previous_effective_under_240s"
    except (OSError, json.JSONDecodeError, TypeError, ValueError):
        pass
    return max(4, min(64, base)), reason


def completed_evidence_ids() -> set[str]:
    """Exclude records already covered by a persisted batch evidence file.

    Some native object gates write batch evidence before the canonical matrix
    projection. Reading only the compact evidence pointer prevents planner
    fast-path loops from selecting the same records again.
    """
    global _PENDING_EVIDENCE_CACHE
    pointers: set[str] = set()
    pending_from_scan: set[str] = set()
    evidence_paths: set[Path] = set()
    for path in (TASK_STATE, CONTINUATION_STATE):
        try:
            payload = json.loads(path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            continue
        candidates = []
        for key in ("batch_checkpoint", "batchCheckpoint", "currentBatchRecovery"):
            value = payload.get(key)
            if isinstance(value, dict) and isinstance(value.get("evidence"), str):
                candidates.append(value["evidence"])
        for candidate in candidates:
            evidence_path = ROOT / candidate
            evidence_paths.add(evidence_path)
    # Canonical manifests may point at family-level runtime evidence. Include
    # those paths, but inspect each record's own status below: a pass_with_
    # boundaries file often contains both exposed and method_evidence_pending
    # records, and the latter must remain schedulable.
    for manifest_name in ("python-api-manifest.json", "object-contract-matrix.json", "workbench-command-matrix.json"):
        try:
            manifest = json.loads((P6 / manifest_name).read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            continue
        for row in manifest.get("records", []) if isinstance(manifest, dict) else []:
            if not isinstance(row, dict):
                continue
            for candidate in row.get("dispositionEvidence", []) if isinstance(row.get("dispositionEvidence"), list) else []:
                if isinstance(candidate, str):
                    evidence_paths.add(ROOT / candidate)
    # A stale manifest can survive an interrupted turn after the checkpoint
    # moves on. Include bounded batch evidence files so planner selection is
    # idempotent even when no compact pointer names the stale batch.
    if P6.exists():
        evidence_paths.update(P6.glob("*module-batch*runtime.json"))
    try:
        cached_payload = json.loads(EVIDENCE_INDEX.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        cached_payload = {}
    cached_entries = cached_payload.get("entries", {}) if isinstance(cached_payload, dict) and cached_payload.get("schema") == "freecad-web-module-batch-evidence-index/v4" else {}
    if not isinstance(cached_entries, dict):
        cached_entries = {}
    next_entries: dict[str, dict[str, object]] = {}
    cache_changed = False
    for evidence_path in sorted(evidence_paths, key=str):
        try:
            stat = evidence_path.stat()
        except OSError:
            continue
        cache_key = str(evidence_path.relative_to(ROOT)) if evidence_path.is_relative_to(ROOT) else str(evidence_path)
        cached = cached_entries.get(cache_key) if isinstance(cached_entries.get(cache_key), dict) else {}
        if (
            cached.get("size") == stat.st_size
            and cached.get("mtime_ns") == stat.st_mtime_ns
            and cached.get("inode") == stat.st_ino
            and isinstance(cached.get("terminalIds"), list)
            and isinstance(cached.get("pendingIds"), list)
        ):
            next_entries[cache_key] = cached
            pointers.update(value for value in cached["terminalIds"] if isinstance(value, str))
            pending_from_scan.update(value for value in cached["pendingIds"] if isinstance(value, str))
            continue
        try:
            evidence = json.loads(evidence_path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            continue
        evidence_state = str(evidence.get("status") or (evidence.get("result") if isinstance(evidence.get("result"), str) else "") or "")
        if not (evidence_state in {"pass", "passed", "pass_with_boundaries", "completed", "blocked_native_probe", "pass_with_restricted_and_pending"} or evidence_state.startswith("pass_with_")):
            continue
        result_payload = evidence.get("result") if isinstance(evidence.get("result"), dict) else {}
        rows = evidence.get("records", []) if isinstance(evidence.get("records"), list) else result_payload.get("records", [])
        top_level_state = str(evidence.get("status") or "")
        is_batch_evidence = "module-batch" in evidence_path.name or "module-batch" in str(evidence.get("schema", ""))
        terminal_ids: list[str] = []
        pending_ids: list[str] = []
        for row in rows:
            if isinstance(row, dict):
                record_id = row.get("apiId") or row.get("id")
                record_state = str(row.get("status") or row.get("result") or "")
                if isinstance(record_id, str) and (
                    top_level_state == "blocked_native_probe"
                    or record_state in TERMINAL_RECORD_STATES
                ):
                    pointers.add(record_id)
                    terminal_ids.append(record_id)
                elif is_batch_evidence and isinstance(record_id, str) and record_state in UNRESOLVED:
                    pending_ids.append(record_id)
                    pending_from_scan.add(record_id)
        next_entries[cache_key] = {
            "size": stat.st_size,
            "mtime_ns": stat.st_mtime_ns,
            "inode": stat.st_ino,
            "terminalIds": sorted(set(terminal_ids)),
            "pendingIds": sorted(set(pending_ids)),
        }
        cache_changed = True
    if set(next_entries) != set(cached_entries):
        cache_changed = True
    if cache_changed:
        EVIDENCE_INDEX.parent.mkdir(parents=True, exist_ok=True)
        temporary = EVIDENCE_INDEX.with_suffix(".tmp")
        temporary.write_text(
            json.dumps({"schema": "freecad-web-module-batch-evidence-index/v4", "entries": next_entries}, ensure_ascii=False, separators=(",", ":")) + "\n",
            encoding="utf-8",
        )
        temporary.replace(EVIDENCE_INDEX)
    _PENDING_EVIDENCE_CACHE = pending_from_scan
    return pointers


def pending_evidence_ids() -> set[str]:
    """Return records with a real runtime attempt but no terminal status yet."""
    if _PENDING_EVIDENCE_CACHE is not None:
        return set(_PENDING_EVIDENCE_CACHE)
    try:
        payload = json.loads(EVIDENCE_INDEX.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return set()
    entries = payload.get("entries", {}) if isinstance(payload, dict) else {}
    pending: set[str] = set()
    for entry in entries.values() if isinstance(entries, dict) else []:
        if isinstance(entry, dict):
            pending.update(value for value in entry.get("pendingIds", []) if isinstance(value, str))
    return pending


def load_records(path: Path, kind: str) -> list[dict]:
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (FileNotFoundError, json.JSONDecodeError):
        return []
    rows = data.get("records", []) if isinstance(data, dict) else []
    result = []
    for row in rows:
        if not isinstance(row, dict):
            continue
        module = row.get("module")
        if not module:
            continue
        if kind == "python":
            disposition = str(row.get("disposition", "planned"))
            if os.environ.get("FREECAD_IMPLEMENT_BOUNDARIES", "0").lower() in {"1", "true", "yes", "on"} and disposition in {"restricted", "nonportable"}:
                disposition = "not_implemented"
            result.append(
                {
                    "id": row.get("id"),
                    "module": module,
                    "kind": row.get("kind", "function"),
                    "name": row.get("name"),
                    "family": row.get("family", "other"),
                    "tier": row.get("tier", "tier-2-pure"),
                    "disposition": disposition,
                    "dispositionEvidence": row.get("dispositionEvidence"),
                    "source": row.get("evidence", {}).get("path") if isinstance(row.get("evidence"), dict) else None,
                    "sourceLine": row.get("evidence", {}).get("line") if isinstance(row.get("evidence"), dict) else None,
                }
            )
        elif kind == "object":
            status = str(row.get("status", "planned"))
            disposition = status if status in {"exposed", "passed", "pass", "completed"} else "not_implemented"
            result.append(
                {
                    "id": row.get("id"),
                    "module": module,
                    "kind": "object",
                    "name": row.get("name"),
                    "family": row.get("adapterFamily", "other"),
                    "tier": "object-contract",
                    "disposition": disposition,
                    "dispositionEvidence": row.get("dispositionEvidence"),
                    "source": (row.get("evidence") or [{}])[0].get("path") if isinstance(row.get("evidence"), list) and row.get("evidence") else None,
                    "sourceLine": (row.get("evidence") or [{}])[0].get("line") if isinstance(row.get("evidence"), list) and row.get("evidence") else None,
                }
            )
        elif kind == "workbench":
            disposition = str(row.get("disposition", "restricted"))
            if os.environ.get("FREECAD_IMPLEMENT_BOUNDARIES", "0").lower() in {"1", "true", "yes", "on"} and disposition == "restricted":
                disposition = "not_implemented"
            result.append(
                {
                    "id": row.get("id"),
                    "module": module,
                    "kind": row.get("kind", "desktop-gui"),
                    "name": row.get("command"),
                    "family": "workbench-command",
                    "tier": "tier-4-restricted",
                    "disposition": disposition,
                    "dispositionEvidence": row.get("dispositionEvidence"),
                    "source": row.get("sourcePath"),
                    "sourceLine": row.get("sourceLine"),
                }
            )
        elif kind == "format":
            status = str(row.get("status", "passed_boundary"))
            result.append(
                {
                    "id": row.get("id"),
                    "module": module,
                    "kind": "format-handler",
                    "name": f"{row.get('direction', 'unknown')}:{row.get('extension', '')}",
                    "family": f"format-{row.get('direction', 'unknown')}",
                    "tier": "format-handler",
                    "disposition": "exposed" if status == "passed" else "not_implemented",
                    "dispositionEvidence": row.get("dispositionEvidence"),
                    "source": row.get("path"),
                    "sourceLine": row.get("line"),
                }
            )
    return [row for row in result if row.get("id") and row.get("name")]


def score(group: tuple[str, str, str], rows: list[dict]) -> tuple:
    module, tier, family = group
    return (
        PREFERRED_MODULES.get(module, 50),
        FAMILY_ORDER.get(family, 20),
        TIER_ORDER.get(tier, 9),
        -sum(row.get("kind") in {"function", "method"} for row in rows),
        -len(rows),
        module,
        tier,
        family,
    )


def compact_row(row: dict) -> list[object]:
    return [row["id"], row["name"], row["disposition"], row.get("source"), row.get("sourceLine")]


def size_without_single_record_tail(total: int, desired: int, minimum: int, cap: int) -> int:
    """Keep a family tail attached to a module batch instead of isolating one API."""
    selected = min(total, desired, cap)
    remaining = total - selected
    if remaining == 1 and selected > minimum:
        selected -= 1
    return max(1, selected)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--module", help="restrict selection to one source module")
    parser.add_argument("--family", help="restrict selection to one adapter family")
    parser.add_argument("--tier", help="restrict selection to one runtime tier")
    parser.add_argument("--size", type=int, default=None, help="maximum records in one batch; omitted means adaptive (default starts at 24)")
    parser.add_argument("--output", type=Path, default=OUTPUT)
    args = parser.parse_args()

    if implementation_first_enabled() and os.environ.get("FREECAD_ALLOW_LEGACY_API_PLANNER", "0").lower() not in {"1", "true", "yes", "on"}:
        payload = {
            "schema": "freecad-web-module-function-batch/v1",
            "status": "no_ready_batch",
            "generatedAt": datetime.now(timezone.utc).isoformat(),
            "reason": "legacy API planner is disabled; use prepare_implementation_business_batch.py",
            "records": [],
            "queue": [],
        }
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(payload, ensure_ascii=False, separators=(",", ":")) + "\n", encoding="utf-8")
        print(json.dumps({"status": "no_ready_batch", "records": 0, "reason": payload["reason"]}, ensure_ascii=False))
        return 0

    records = []
    records.extend(load_records(P6 / "python-api-manifest.json", "python"))
    records.extend(load_records(P6 / "object-contract-matrix.json", "object"))
    records.extend(load_records(P6 / "workbench-command-matrix.json", "workbench"))
    records.extend(load_records(P6 / "format-coverage-kernel.json", "format"))
    if implementation_first_enabled():
        # Canonical exposed/passed/native_handler rows are already excluded by
        # disposition.  For unresolved rows, implementation receipts are the
        # only reusable state; historical gate files are diagnostics and must
        # not be rescanned or mistaken for completed business behavior.
        covered_ids = all_implementation_ready_ids()
        pending_ids: set[str] = set()
    else:
        covered_ids = completed_evidence_ids()
        pending_ids = pending_evidence_ids() if os.environ.get("FREECAD_RETRY_EVIDENCED", "0").lower() not in {"1", "true", "yes", "on"} else set()
    groups: dict[tuple[str, str, str], list[dict]] = defaultdict(list)
    for row in records:
        if row["disposition"] not in UNRESOLVED:
            continue
        if row.get("id") in covered_ids:
            continue
        if row.get("id") in pending_ids:
            continue
        if args.module and row["module"] != args.module:
            continue
        if args.family and row["family"] != args.family:
            continue
        if args.tier and row["tier"] != args.tier:
            continue
        groups[(row["module"], row["tier"], row["family"])].append(row)
    for rows in groups.values():
        rows.sort(key=lambda row: (row.get("source") or "", row.get("sourceLine") or 0, row["name"], row["id"]))
    if not groups:
        payload = {
            "schema": "freecad-web-module-function-batch/v1",
            "status": "no_ready_batch",
            "generatedAt": datetime.now(timezone.utc).isoformat(),
            "records": [],
            "queue": [],
        }
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(payload, ensure_ascii=False, separators=(",", ":")) + "\n", encoding="utf-8")
        print(json.dumps({"status": payload["status"], "records": 0}, ensure_ascii=False))
        # An empty queue is a valid terminal planner result. Returning a
        # non-zero status makes the scheduler treat an exhausted module as a
        # command failure even though the structured payload is complete.
        return 0

    ordered = sorted(groups.items(), key=lambda item: score(item[0], item[1]))
    min_records = 2
    try:
        checkpoint = json.loads((ROOT / "docs/current/continuation-state.json").read_text(encoding="utf-8"))
        min_records = max(1, int(checkpoint.get("batchPolicy", {}).get("minRecords", 2)))
    except (OSError, json.JSONDecodeError, TypeError, ValueError):
        pass
    # Prefer a full-size group whenever one exists. Singleton groups are only
    # selected after all remaining groups are smaller than the policy minimum;
    # this prevents the scheduler from repeatedly stopping on a one-record
    # tail while larger work is still available.
    eligible = [item for item in ordered if len(item[1]) >= min_records]
    singleton_fallback = not eligible
    selection_pool = eligible or ordered
    (module, tier, family), selected_rows = selection_pool[0]
    size, size_reason = adaptive_target_size(args.size, tier, module, family)
    if tier == "object-contract":
        family_cap = 16
    elif family in {"pure-function", "geometry-value"}:
        family_cap = 32
    else:
        family_cap = 24
    effective_size = min(size, family_cap)
    # Part generic-object fixtures construct heavyweight OCCT 2D/3D wrapper
    # objects. Keep their batch below the 32 MiB fast-dev initial heap so one
    # malformed or unusually large record cannot exhaust the whole probe.
    if module == "Part" and family == "generic-object":
        effective_size = min(effective_size, 12)
    effective_size = size_without_single_record_tail(
        len(selected_rows), effective_size, min_records, family_cap
    )
    selected_rows = selected_rows[:effective_size]
    counts = Counter(row["disposition"] for row in selected_rows)
    identity = "|".join([module, tier, family, *(row["id"] for row in selected_rows)])
    batch_id = hashlib.sha256(identity.encode("utf-8")).hexdigest()[:24]
    queue = []
    for (q_module, q_tier, q_family), q_rows in ordered[:12]:
        queue.append(
            {
                "module": q_module,
                "tier": q_tier,
                "family": q_family,
                "unresolved": len(q_rows),
                "nextSize": min(effective_size, len(q_rows)),
            }
        )
    state = json.loads(STATE.read_text(encoding="utf-8"))
    module_groups = [
        {"tier": q_tier, "family": q_family, "remaining": len(q_rows)}
        for (q_module, q_tier, q_family), q_rows in ordered
        if q_module == module
    ]
    module_records = [row for row in records if row.get("module") == module]
    module_reused = sum(
        row.get("disposition") not in UNRESOLVED or str(row.get("id")) in covered_ids
        for row in module_records
    )
    payload = {
        "schema": "freecad-web-module-function-batch/v1",
        "status": "ready",
        "generatedAt": datetime.now(timezone.utc).isoformat(),
        "taskId": state.get("active_task"),
        "batchId": batch_id,
        "strategy": "module-business-first-deferred-validation",
        "businessBundle": {
            "module": module,
            "coverageRecords": sum(row["remaining"] for row in module_groups),
            "families": module_groups,
            "selectedRecordsAreCoverageAnchors": True,
            "implementationUnit": "shared-module-dispatcher",
            "reuse": {
                "moduleInventoryRecords": len(module_records),
                "moduleAlreadyImplementedRecords": module_reused,
                "moduleRemainingRecords": len(module_records) - module_reused,
                "globalImplementationReceipts": len(all_implementation_ready_ids()) if implementation_first_enabled() else 0,
            },
        },
        "targetSize": effective_size,
        "targetSizeReason": size_reason,
        "wasmPolicy": {
            "mode": "independent-module-artifact",
            "dynamicLoad": True,
            "mergeRequired": False,
            "artifactFields": ["module", "worker", "revision", "runtimeUrl", "wasmUrl", "dataUrl", "assetBaseUrl"],
        },
        "minimumRecords": min_records,
        "singletonFallback": singleton_fallback and len(selected_rows) < min_records,
        "selected": {"module": module, "tier": tier, "family": family, "recordCount": len(selected_rows), "dispositionCounts": dict(counts)},
        "fields": ["id", "name", "disposition", "source", "sourceLine"],
        "records": [compact_row(row) for row in selected_rows],
        "queue": queue,
        "boundaries": [
            "The batch planner never changes dispositions or task state.",
            "Selected API/object rows are coverage anchors for one module business implementation, never individual implementation tasks.",
            "One module dispatcher/family adapter is reused for the entire batch; no API-specific Worker, page, build or gate is created.",
            "WASM artifacts remain split by functional module and are loaded on demand; no binary merge is required.",
            "Implementation mode performs target build/sync/receipt only; validation remains globally deferred.",
        ],
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, ensure_ascii=False, separators=(",", ":")) + "\n", encoding="utf-8")
    print(json.dumps({"status": payload["status"], "batchId": batch_id, "module": module, "tier": tier, "family": family, "records": len(selected_rows), "queue": queue[:4]}, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())