#!/usr/bin/env python3
"""Generate persisted isolated burst scenarios and mapped SCADA observations."""
from __future__ import annotations

import argparse
import asyncio
import os
import sys
from datetime import datetime, timezone
from pathlib import Path


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--design", type=Path, default=Path("resources/benchmarks/tjwater2-v2-burst-balanced.json"))
    parser.add_argument("--target", default="tjw_bench_tjwater2_20261105")
    parser.add_argument("--state-dir", type=Path, default=Path("data/benchmark-generation/tjwater2-20261105"))
    parser.add_argument("--apply", action="store_true")
    parser.add_argument("--prepare-only", action="store_true")
    args = parser.parse_args()
    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
    os.umask(0o077)
    from psycopg import connect
    from psycopg.conninfo import make_conninfo
    from psycopg.rows import dict_row
    from app.core.config import settings
    from scripts.benchmark.calibration import save_json, digest
    from scripts.benchmark.generation import (
        apply_project_runtime_migrations,
        create_empty_database,
        copy_project_inputs,
        make_schedule,
        register_benchmark_project,
    )
    from scripts.benchmark.simulation_design import load_simulation_design
    from scripts.check_burst_simulation_design import resolve_project, readonly, validate_database

    design = load_simulation_design(args.design)
    schedule = make_schedule(design)
    if not args.apply:
        print(f"dry run: {len(schedule)} cases, target={args.target}, dates={schedule[0]['history_start']}..{schedule[-1]['reserved_end']}")
        return 0
    folder = args.state_dir.resolve()
    folder.mkdir(parents=True, exist_ok=True, mode=0o700)
    state_path = folder / "state.json"
    import json
    state = json.loads(state_path.read_text()) if state_path.exists() else {
        "version": "burst-dataset-generation@1", "created_at": datetime.now(timezone.utc).isoformat(),
        "target": args.target, "design_sha256": digest(args.design), "schedule": schedule,
        "stages": {}, "cases": {}, "source_effect": "read-only", "status": "preparing"}
    if state["target"] != args.target or state["design_sha256"] != digest(args.design):
        raise ValueError("state belongs to a different target or design")
    save_json(state_path, state)
    biz, iot = asyncio.run(resolve_project(design))
    with readonly(biz) as connection:
        validate_database(connection, design, template=False)
    target_biz, target_iot = make_conninfo(biz, dbname=args.target), make_conninfo(iot, dbname=args.target)
    for role, source, template in (("biz", biz, settings.WNDB_SCHEMA_TEMPLATE_DB_NAME),
                                   ("iot", iot, settings.TIMESCALEDB_SCHEMA_TEMPLATE_DB_NAME)):
        if not state["stages"].get(f"{role}_created"):
            create_empty_database(source, args.target, template)
            state["stages"][f"{role}_created"] = True
            save_json(state_path, state)
            print(f"created isolated {role} database: {args.target}", flush=True)
    if not state["stages"].get("model_copied"):
        state["model_counts"] = copy_project_inputs(biz, target_biz)
        state["stages"]["model_copied"] = True
        save_json(state_path, state)
        print("copied model and device bindings from a read-only source snapshot", flush=True)
    required_migrations = {
        "025_mcp_query_contexts.sql",
        "026_calculation_jobs.sql",
        "027_calculation_job_lifecycle.sql",
        "030_remove_calculation_queue_expiry.sql",
    }
    if not required_migrations.issubset(state.get("runtime_migrations", [])):
        state["runtime_migrations"] = apply_project_runtime_migrations(
            target_biz, Path(__file__).resolve().parents[1]
        )
        state["stages"]["runtime_migrations_applied"] = True
        save_json(state_path, state)
        print("installed isolated MCP and calculation queue storage", flush=True)
    if not state["stages"].get("model_template_created"):
        from app.native.wndb.core.project_templates import (
            create_project_model_template,
            ensure_replication_worker_capacity,
        )
        from app.native.wndb.core.projects import have_project

        template = f"{args.target}_template"
        if not have_project(template):
            try:
                ensure_replication_worker_capacity()
            except RuntimeError as exc:
                if "no reserved logical-replication worker capacity" not in str(exc):
                    raise
                state["model_template_mode"] = "snapshot_per_run"
            else:
                create_project_model_template(args.target)
        if have_project(template):
            state["model_template"] = template
            state["model_template_mode"] = "logical_replication"
        state["stages"]["model_template_created"] = True
        save_json(state_path, state)
        print(f"configured hydraulic model isolation: {state['model_template_mode']}", flush=True)
    if not state["stages"].get("metadata_registered"):
        project = asyncio.run(register_benchmark_project(
            source_project_id=str(design.project_id),
            target=args.target,
            target_biz_dsn=target_biz,
            target_iot_dsn=target_iot,
        ))
        state["project_id"] = str(project.id)
        state["stages"]["metadata_registered"] = True
        save_json(state_path, state)
        print(f"registered isolated benchmark project: {project.id}", flush=True)
    # Inspect isolated schema readiness before invoking backend persistence.
    with connect(target_biz, row_factory=dict_row) as connection:
        print("isolated BizDB analysis schema", connection.execute(
            "SELECT column_name FROM information_schema.columns WHERE table_schema='analysis' AND table_name='runs' ORDER BY ordinal_position").fetchall(), flush=True)
    state["status"] = "prepared"
    save_json(state_path, state)
    print(f"isolated preparation complete: {state_path}", flush=True)
    return 0


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