#!/usr/bin/env python3
"""Offline, bounded pilot for context-aware procedural graphs and verifiable completion.

This is an engineering-control experiment, not a claim of new science or real-world
agent improvement. It uses deterministic synthetic workflows, standard-library code,
and no network, model-generated code, production APIs, or external actions.
"""
from __future__ import annotations

import argparse
import json
import os
import tempfile
from collections import defaultdict
from pathlib import Path
from typing import Any

TASK_PATHS: dict[str, list[list[str]]] = {
    "network_timeout": [
        ["observe", "classify", "check_egress", "probe_endpoint", "verify"],
        ["observe", "classify", "check_egress", "use_relay", "verify"],
    ],
    "disk_pressure": [
        ["observe", "classify", "measure_disk", "prune_verified_images", "verify"],
        ["observe", "classify", "measure_disk", "remove_temp_files", "verify"],
    ],
    "deployment_failure": [
        ["observe", "classify", "inspect_build", "rollback_candidate", "verify"],
        ["observe", "classify", "inspect_build", "retry_build", "verify"],
    ],
}
NEGATIVE_ACTION = "repeat_diagnosis"
ALPHA = 1.0
BETA = 1.0


def make_dataset() -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, dict[str, set[str]]], list[str]]:
    """Create deterministic train/held-out workflow traces plus the evaluator's allowed edges."""
    train: list[dict[str, Any]] = []
    test: list[dict[str, Any]] = []
    allowed: dict[str, dict[str, set[str]]] = defaultdict(lambda: defaultdict(set))
    actions: set[str] = {NEGATIVE_ACTION}
    for context, paths in TASK_PATHS.items():
        for path_index, trace in enumerate(paths):
            for source, target in zip(trace, trace[1:]):
                allowed[context][source].add(target)
                actions.update((source, target))
            # Duplicate only independent synthetic instances; hold-out IDs never train the graph.
            for replica in range(6):
                item = {"context": context, "trace": trace[:], "case_id": f"{context}:{path_index}:{replica}"}
                (train if replica < 4 else test).append(item)
        # A reviewed simulated failure records an unproductive transition, never a real action.
        for replica in range(2):
            train.append({
                "context": context,
                "trace": ["observe", "classify", NEGATIVE_ACTION],
                "case_id": f"{context}:negative:{replica}",
                "success": False,
                "failure_index": 1,
            })
    return train, test, allowed, sorted(actions)


def fit_graph(records: list[dict[str, Any]], context_aware: bool) -> dict[tuple[str, str, str], dict[str, int]]:
    """Count successful and failed transitions; failed traces penalize only their failed edge."""
    counts: dict[tuple[str, str, str], dict[str, int]] = defaultdict(lambda: {"success": 0, "failure": 0})
    for record in records:
        context = str(record["context"]) if context_aware else "*"
        trace = record["trace"]
        success = record.get("success", True) is True
        failure_index = int(record.get("failure_index", -1))
        for index, (source, target) in enumerate(zip(trace, trace[1:])):
            if success or index < failure_index:
                counts[(context, source, target)]["success"] += 1
            elif index == failure_index:
                counts[(context, source, target)]["failure"] += 1
    return dict(counts)


def edge_probability(counts: dict[str, int]) -> float:
    return (counts.get("success", 0) + ALPHA) / (
        counts.get("success", 0) + counts.get("failure", 0) + ALPHA + BETA
    )


def predict_next(
    graph: dict[tuple[str, str, str], dict[str, int]],
    context: str,
    source: str,
    actions: list[str],
    context_aware: bool,
) -> tuple[str | None, str]:
    """Pick a supported next procedure; fail closed when a context-aware graph has no context evidence."""
    scope = context if context_aware else "*"
    if context_aware and not any(key[0] == scope and key[1] == source for key in graph):
        return None, "no_context_specific_evidence"
    outgoing = []
    for target in actions:
        counts = graph.get((scope, source, target))
        if counts is None:
            continue
        outgoing.append((edge_probability(counts), counts.get("success", 0), target))
    if not outgoing:
        return None, "no_transition_evidence"
    outgoing.sort(key=lambda row: (-row[0], -row[1], row[2]))
    probability, successes, target = outgoing[0]
    if successes == 0 or probability < 0.60:
        return None, "insufficient_success_evidence"
    return target, "evidence_selected"


def validate_success_claim(gate_results: list[dict[str, Any]]) -> dict[str, Any]:
    """Only permit terminal success with at least one passed, evidence-linked gate."""
    if not gate_results:
        raise ValueError("success_requires_at_least_one_gate")
    for gate in gate_results:
        if gate.get("passed") is not True:
            raise ValueError("success_requires_all_gates_to_pass")
        if not isinstance(gate.get("evidence_ref"), str) or not gate["evidence_ref"].strip():
            raise ValueError("success_requires_gate_evidence_reference")
    return {"status": "succeeded", "gates": gate_results}


def atomic_state_write(path: Path, state: dict[str, Any]) -> None:
    """Durably replace one JSON state file; safe to call from a one-tick worker."""
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temp_name = tempfile.mkstemp(prefix=path.name + ".tmp-", dir=path.parent)
    try:
        os.fchmod(fd, 0o600)
        with os.fdopen(fd, "w", encoding="utf-8") as stream:
            json.dump(state, stream, sort_keys=True, ensure_ascii=False)
            stream.write("\n")
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temp_name, path)
        directory_fd = os.open(path.parent, os.O_RDONLY)
        try:
            os.fsync(directory_fd)
        finally:
            os.close(directory_fd)
    finally:
        try:
            os.unlink(temp_name)
        except FileNotFoundError:
            pass


def evaluate(
    graph: dict[tuple[str, str, str], dict[str, int]],
    cases: list[dict[str, Any]],
    allowed: dict[str, dict[str, set[str]]],
    actions: list[str],
    context_aware: bool,
) -> dict[str, Any]:
    totals: dict[str, list[int]] = defaultdict(lambda: [0, 0])
    correct = 0
    decisions = 0
    for case in cases:
        context = case["context"]
        seen_sources: set[str] = set()
        for source, _observed_target in zip(case["trace"], case["trace"][1:]):
            if source in seen_sources:
                continue
            seen_sources.add(source)
            valid_targets = allowed[context].get(source, set())
            if not valid_targets:
                continue
            prediction, _reason = predict_next(graph, context, source, actions, context_aware)
            decisions += 1
            totals[context][1] += 1
            if prediction in valid_targets:
                correct += 1
                totals[context][0] += 1
    return {
        "correct": correct,
        "decisions": decisions,
        "accuracy": round(correct / decisions, 6) if decisions else 0.0,
        "per_context": {
            key: {
                "correct": values[0],
                "decisions": values[1],
                "accuracy": round(values[0] / values[1], 6) if values[1] else 0.0,
            }
            for key, values in sorted(totals.items())
        },
    }


def run_experiment(cycle: int = 0, include_literature: bool = False) -> dict[str, Any]:
    train, test, allowed, actions = make_dataset()
    baseline_graph = fit_graph(train, context_aware=False)
    candidate_graph = fit_graph(train, context_aware=True)
    baseline = evaluate(baseline_graph, test, allowed, actions, context_aware=False)
    candidate = evaluate(candidate_graph, test, allowed, actions, context_aware=True)

    unknown_prediction, unknown_reason = predict_next(
        candidate_graph, "unseen_context", "classify", actions, context_aware=True
    )
    rejected_empty = False
    rejected_failed = False
    rejected_missing_reference = False
    passed_gate_allowed = False
    try:
        validate_success_claim([])
    except ValueError:
        rejected_empty = True
    try:
        validate_success_claim([{"name": "smoke-test", "passed": False, "evidence_ref": "test://1"}])
    except ValueError:
        rejected_failed = True
    try:
        validate_success_claim([{"name": "smoke-test", "passed": True}])
    except ValueError:
        rejected_missing_reference = True
    try:
        allowed_claim = validate_success_claim([
            {"name": "smoke-test", "passed": True, "evidence_ref": "sha256:" + ("a" * 64)}
        ])
        passed_gate_allowed = allowed_claim["status"] == "succeeded"
    except ValueError:
        passed_gate_allowed = False

    no_regression = all(
        candidate["per_context"].get(context, {}).get("accuracy", 0.0)
        >= baseline["per_context"].get(context, {}).get("accuracy", 0.0)
        for context in TASK_PATHS
    )
    accepted = (
        candidate["accuracy"] > baseline["accuracy"]
        and no_regression
        and unknown_prediction is None
    )
    validation = (
        accepted
        and rejected_empty
        and rejected_failed
        and rejected_missing_reference
        and passed_gate_allowed
    )
    return {
        "task": "procedural_graph_and_verifiable_completion_pilot",
        "status": "ok" if validation else "failed",
        "cycle": cycle,
        "method": "Deterministic synthetic workflows; context-agnostic baseline versus context-aware transition graph trained only on training traces; failed transitions penalized; held-out decisions checked against a fixed evaluator specification.",
        "dataset": {
            "train_records": len(train),
            "held_out_records": len(test),
            "contexts": sorted(TASK_PATHS),
            "actions": len(actions),
            "external_calls": 0,
            "model_generated_code_executed": False,
        },
        "baseline": baseline,
        "candidate": candidate,
        "absolute_gain_percentage_points": round((candidate["accuracy"] - baseline["accuracy"]) * 100, 3),
        "candidate_accepted_by_gate": accepted,
        "no_per_context_regression": no_regression,
        "unseen_context_abstention": {
            "passed": unknown_prediction is None,
            "prediction": unknown_prediction,
            "reason": unknown_reason,
        },
        "completion_gate_tests": {
            "empty_gates_rejected": rejected_empty,
            "failed_gate_rejected": rejected_failed,
            "missing_evidence_reference_rejected": rejected_missing_reference,
            "passed_evidence_linked_gate_accepted": passed_gate_allowed,
        },
        "validation_passed": validation,
        "claim": (
            f"On this deterministic synthetic workflow benchmark only, the context-aware procedural graph "
            f"scored {candidate['accuracy']:.1%} versus {baseline['accuracy']:.1%} for the context-agnostic "
            f"baseline ({(candidate['accuracy'] - baseline['accuracy']) * 100:.1f} percentage points). "
            "The pilot also rejected unsupported success claims and abstained on an unseen context. "
            "This validates the pilot implementation's control behavior only; it does not establish improved "
            "real-world LLM-agent performance, general intelligence, or a new scientific discovery."
            if validation else
            "Pilot validation failed; the result must not be promoted to verified project knowledge."
        ),
        "novelty_status": "not a novelty claim; engineering pilot inspired by prior published work",
        "model_weights_updated": False,
        "production_graph_updated": False,
        "literature_lookup_requested": include_literature,
    }


def self_test() -> dict[str, Any]:
    result = run_experiment()
    assert result["validation_passed"], result
    assert result["candidate"]["accuracy"] > result["baseline"]["accuracy"], result
    assert result["unseen_context_abstention"]["passed"], result
    assert all(result["completion_gate_tests"].values()), result
    with tempfile.TemporaryDirectory(prefix="darkoza-procedural-graph-test-") as temp_dir:
        path = Path(temp_dir) / "state.json"
        state = {"task_id": "test-1", "status": "verifying", "tick": 3}
        atomic_state_write(path, state)
        assert json.loads(path.read_text(encoding="utf-8")) == state
        assert (path.stat().st_mode & 0o777) == 0o600
    return {"status": "passed", "assertions": 6, "experiment": result}


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true")
    parser.add_argument("--cycle", type=int, default=0)
    args = parser.parse_args()
    result = self_test() if args.self_test else run_experiment(cycle=args.cycle)
    print(json.dumps(result, ensure_ascii=False, sort_keys=True))
    return 0 if result.get("status") == "passed" or result.get("validation_passed") is True else 1


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