#!/usr/bin/env python3
"""Generate, validate, and record independent implementations of a spec function."""

from __future__ import annotations

import argparse
import datetime as dt
import hashlib
import json
import os
from pathlib import Path
import shlex
import subprocess
import sys
import tempfile
import tomllib
from typing import Any


DEFAULT_AGENT_COMMAND = (
    "codex exec --ephemeral --sandbox read-only "
    "--output-last-message {output} -"
)

INTERCHANGEABILITY_RULE = (
    "An implementation may read or write external resources, but it must not retain hidden "
    "implementation-local persistent state that another spec function or future invocation "
    "depends upon. State shared across functions or calls must be explicit in arguments, "
    "results, shared context, or a declared external resource."
)

DIVERSITY_DIMENSIONS = {
    "algorithm": [
        "direct domain algorithm using explicit arithmetic or comparisons",
        "table- or rule-driven decision structure",
        "small composable transformations with sharply separated responsibilities",
        "derive the result from invariants or algebraic relationships",
        "explicit state-transition or decision model",
    ],
    "control_flow": [
        "guard clauses and early returns",
        "exhaustive pattern matching",
        "clearly separated processing phases",
        "iterator or pipeline composition",
        "one centralized decision table",
    ],
    "data_representation": [
        "domain-specific enums or structs",
        "compact primitive values",
        "immutable intermediate values",
        "named constants and lookup data",
        "type-directed wrapper values",
    ],
    "validation": [
        "validate all inputs before transformation",
        "validate at the public boundary, then operate on trusted values",
        "validate each value immediately before it is used",
        "make invalid states explicit and handle cases exhaustively",
        "separate validation, transformation, and output construction",
    ],
    "engineering_priority": [
        "auditability and obvious correctness",
        "bounded resource use",
        "minimal allocation and copying",
        "simple, independently testable control flow",
        "strict edge-case handling",
    ],
    "dependency_policy": [
        "standard library only",
        "existing project dependencies only; add nothing new",
        "prefer language primitives and project-owned helpers",
    ],
}


def diversity_brief(function: str, candidate: int) -> dict[str, str]:
    digest = hashlib.sha256(function.encode()).digest()
    strides = [1, 2, 3, 4, 1, 2]
    brief = {}
    for dimension_index, (name, options) in enumerate(DIVERSITY_DIMENSIONS.items()):
        option_index = (digest[dimension_index] + (candidate - 1) * strides[dimension_index]) % len(options)
        brief[name] = options[option_index]
    return brief


def load_contract(path: Path, name: str) -> tuple[dict[str, Any], dict[str, Any]]:
    with path.open("rb") as handle:
        document = tomllib.load(handle)
    for function in document.get("function", []):
        if function.get("name") == name:
            return document, function
    available = ", ".join(f["name"] for f in document.get("function", []))
    raise SystemExit(f"unknown spec function {name!r}; available: {available}")


def build_prompt(document: dict[str, Any], contract: dict[str, Any], context: str, index: int) -> str:
    prior = "No earlier candidates are shown: choose your own algorithm and decomposition."
    brief = diversity_brief(contract["name"], index)
    return f"""You are producing candidate {index} for a swappable implementation.

Application: {document['application']} {document['release']}
External references: {', '.join(document.get('external_references', []))}
Mandatory Polyform interchangeability rule: {INTERCHANGEABILITY_RULE}

Diversity brief for this candidate:
{json.dumps(brief, indent=2)}

Contract (JSON):
{json.dumps(contract, indent=2)}

Relevant project context:
{context}

{prior}

Return ONLY the complete Rust source for the candidate module. Use private helpers freely.
Do not change the contract, tests, public types, or other modules. Favor a meaningfully
independent control flow or primitive over cosmetic renaming. Follow the diversity brief as a
strong design direction and do not copy the structure or helper decomposition of existing
variants. If a dimension does not fit the contract, preserve its underlying diversity goal.
Do not add hidden persistent state.
"""


def run(command: list[str], cwd: Path, *, stdin: str | None = None) -> subprocess.CompletedProcess[str]:
    return subprocess.run(command, cwd=cwd, input=stdin, text=True, capture_output=True, check=False)


def invoke_agent(template: str, prompt: str, output: Path, cwd: Path) -> tuple[str, str]:
    prompt_path = output.with_suffix(".prompt.txt")
    prompt_path.write_text(prompt)
    command = [part.format(prompt=str(prompt_path), output=str(output)) for part in shlex.split(template)]
    result = run(command, cwd, stdin=prompt)
    if result.returncode:
        raise RuntimeError(f"agent command failed ({result.returncode}):\n{result.stderr}")
    if output.exists():
        source = output.read_text()
    else:
        source = result.stdout
        output.write_text(source)
    return source, " ".join(command)


def validate(repo: Path, staged: Path, test_filter: str, regression_command: str) -> list[dict[str, Any]]:
    checks = [
        ["cargo", "fmt", "--", "--check"],
        ["cargo", "check", "--all-targets"],
        ["cargo", "test", test_filter],
        shlex.split(regression_command),
    ]
    outcomes = []
    for command in checks:
        result = run(command, repo)
        outcomes.append({"command": command, "returncode": result.returncode,
                         "stdout": result.stdout[-4000:], "stderr": result.stderr[-4000:]})
        if result.returncode:
            break
    return outcomes


def append_jsonl(path: Path, record: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record, sort_keys=True) + "\n")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--contract", type=Path, default=Path("specs/zipdiv.toml"))
    parser.add_argument("--function", required=True)
    parser.add_argument("--count", type=int, default=1)
    parser.add_argument(
        "--agent-command",
        default=os.environ.get("POLYFORM_AGENT_COMMAND", DEFAULT_AGENT_COMMAND),
        help=(
            "Command template; {prompt} and {output} are replaced. Prompt is also stdin. "
            "Defaults to non-interactive Codex CLI."
        ),
    )
    parser.add_argument("--agent", default=os.environ.get("POLYFORM_AGENT", "codex"))
    parser.add_argument("--model", default=os.environ.get("POLYFORM_MODEL", "configured-default"))
    parser.add_argument("--context", type=Path, action="append", default=[])
    parser.add_argument("--rust-root", type=Path, default=Path("rust"))
    parser.add_argument("--candidate-dir", type=Path, default=Path("rust/src/generated"))
    parser.add_argument("--test-filter", default="contract")
    parser.add_argument("--regression-command", default="cargo test --test regression")
    parser.add_argument("--provenance", type=Path, default=Path("artifacts/provenance.jsonl"))
    parser.add_argument("--dry-run", action="store_true", help="Print prompts without invoking an agent")
    args = parser.parse_args()
    if args.count < 1:
        parser.error("--count must be positive")
    document, contract = load_contract(args.contract, args.function)
    context_parts = []
    for path in args.context:
        context_parts.append(f"--- {path} ---\n{path.read_text()}")
    context = "\n".join(context_parts) or "No additional context supplied."

    accepted = 0
    for index in range(1, args.count + 1):
        prompt = build_prompt(document, contract, context, index)
        if args.dry_run:
            print(prompt)
            continue
        with tempfile.TemporaryDirectory(prefix="polyform-candidate-") as tmp:
            staged = Path(tmp) / f"{args.function}_candidate_{index}.rs"
            source, invocation = invoke_agent(args.agent_command, prompt, staged, Path.cwd())
            digest = hashlib.sha256(source.encode()).hexdigest()
            destination = args.candidate_dir / f"{args.function}_{digest[:12]}.rs"
            destination.parent.mkdir(parents=True, exist_ok=True)
            destination.write_text(source)
            outcomes = validate(args.rust_root, destination, args.test_filter, args.regression_command)
            passed = len(outcomes) == 4 and all(item["returncode"] == 0 for item in outcomes)
            record = {
                "schema_version": 1,
                "function": args.function,
                "agent": args.agent,
                "model": args.model,
                "generated_at": dt.datetime.now(dt.timezone.utc).isoformat(),
                "sha256": digest,
                "prompt": prompt,
                "diversity_brief": diversity_brief(args.function, index),
                "invocation": invocation,
                "candidate": str(destination),
                "accepted": passed,
                "checks": outcomes,
            }
            append_jsonl(args.provenance, record)
            if passed:
                accepted += 1
                print(f"accepted {destination} ({digest[:12]})")
            else:
                destination.unlink(missing_ok=True)
                print(f"rejected candidate {index}; see {args.provenance}", file=sys.stderr)
    return 0 if args.dry_run or accepted == args.count else 1


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