"""python -m app.simulation.native_codegen INPUT.xml|json --output-dir DIR"""
from __future__ import annotations

import argparse
from dataclasses import replace
from hashlib import sha256
import json
from pathlib import Path
import shutil
import statistics
import time

from app.main import compile_system_xml_network
from app.simulation.backends import simulation_config
from .build import build_native
from .compiler import compile_native_program
from .input import load_input
from .runner import execute_native


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", type=Path)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--method", choices=("RK45", "BDF"))
    parser.add_argument("--max-step", type=float)
    parser.add_argument("--rtol", type=float)
    parser.add_argument("--runs", type=int, default=3)
    parser.add_argument("--timeout", type=float, default=300)
    parser.add_argument("--solve-only", action="store_true")
    args = parser.parse_args()
    if args.runs < 1:
        parser.error("--runs must be positive")
    out = args.output_dir.resolve()
    out.mkdir(parents=True, exist_ok=True)
    if (out / "summary.json").exists():
        parser.error("Output directory already contains a completed run; choose a new directory.")
    started = time.perf_counter()
    xml, document = load_input(args.input)
    program = compile_native_program(compile_system_xml_network(document))
    config = simulation_config(document.simulation)
    for name in ("method", "max_step", "rtol"):
        if getattr(args, name) is not None:
            config = replace(config, **{name: getattr(args, name)})
    preparation = time.perf_counter()-started
    (out / "input.xml").write_bytes(xml)
    build = build_native(program)
    package = out / "program"
    package.mkdir(exist_ok=True)
    for name in (*build.manifest["artifacts"], "manifest.json"):
        shutil.copy2(build.executable.parent / name, package / name)
    (out / "model-manifest.json").write_text(json.dumps(build.manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    rows = []
    for i in range(args.runs+1):
        data = execute_native(build, config, document.simulation.sample_step,
                              run_dir=out / ("warmup" if i == 0 else f"run-{i}"),
                              record_samples=not args.solve_only, timeout=args.timeout)
        row = {k: v for k, v in data.items() if k not in ("series", "final", "finalState")}
        row["run"] = "warmup" if i == 0 else i
        rows.append(row)
        print(json.dumps(row, ensure_ascii=False), flush=True)
        if not data["success"]:
            break
    measured = rows[1:]
    summary = {
        "source": str(args.input.resolve()), "sourceSha256": sha256(args.input.read_bytes()).hexdigest(),
        "xmlSha256": sha256(xml).hexdigest(), "executable": str(package / build.executable.name),
        "cachedExecutable": str(build.executable),
        "buildKey": build.manifest["buildKey"], "buildSeconds": build.seconds,
        "preparationSeconds": preparation, "cacheHit": build.cache_hit,
        "settings": vars(config), "sampleStep": document.simulation.sample_step,
        "recordSamples": not args.solve_only, "runs": rows,
        "success": len(measured) == args.runs and all(row["success"] for row in rows),
        "medianSolveSeconds": statistics.median(row["solveSeconds"] for row in measured) if measured else None,
    }
    (out / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps({k: v for k, v in summary.items() if k != "runs"}, ensure_ascii=False))
    return 0 if summary["success"] else 1


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