#!/usr/bin/env python3
from __future__ import annotations

import argparse
import copy
import hashlib
import json
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

SCRIPT_DIR = Path(__file__).resolve().parent
sys.dont_write_bytecode = True
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

from auto_conform_benchmark import (
    ANALYSIS_VERSION,
    AutoConformFailure,
    atomic_write_json,
    benchmark_summary,
    load_schema,
    merge_analysis_diagnostics,
    new_analysis_diagnostics,
    progress,
    validate_contract,
    write_alternative_table,
    write_placement_table,
)


PLAN_FORMAT = "automation-agent-auto-conform-analysis-plan/v1"


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def batch_plan(shots: list[dict[str, Any]], batch_size: int) -> list[dict[str, Any]]:
    batches: list[dict[str, Any]] = []
    for batch_index, start_index in enumerate(range(0, len(shots), batch_size)):
        end_index = min(len(shots), start_index + batch_size)
        batches.append(
            {
                "batchIndex": batch_index,
                "startIndex": start_index,
                "shotCount": end_index - start_index,
                "firstShotNumber": start_index + 1,
                "lastShotNumber": end_index,
                "shotIds": [str(shot["shotId"]) for shot in shots[start_index:end_index]],
                "fileName": f"analysis-batch-{batch_index}.json",
            }
        )
    return batches


def validate_unique_shot_ids(shots: list[dict[str, Any]]) -> None:
    shot_ids = [str(shot["shotId"]) for shot in shots]
    if len(set(shot_ids)) != len(shot_ids):
        raise AutoConformFailure(
            "analysis-plan",
            "DUPLICATE_SHOT_ID",
            "analysis planning requires unique shot IDs",
        )


def plan_analysis(args: argparse.Namespace) -> int:
    manifest_path = Path(args.manifest).resolve()
    schema = load_schema(Path(args.schema).resolve())
    document = json.loads(manifest_path.read_text(encoding="utf-8"))
    validate_contract(document, schema)
    shots = document["shots"]
    if not shots:
        raise AutoConformFailure(
            "analysis-plan",
            "NO_SHOTS",
            "analysis planning requires at least one reference shot",
        )
    validate_unique_shot_ids(shots)
    batches = batch_plan(shots, args.batch_size)
    parts_dir = Path(args.parts_dir).resolve()
    tasks_path = Path(args.tasks_output).resolve()
    plan_path = Path(args.plan_output).resolve()
    parts_dir.mkdir(parents=True, exist_ok=True)
    tasks_path.parent.mkdir(parents=True, exist_ok=True)
    task_lines = [
        "\t".join(
            [
                str(batch["batchIndex"]),
                str(batch["batchIndex"] + 1),
                str(len(batches)),
                str(batch["startIndex"]),
                str(batch["shotCount"]),
                str(batch["firstShotNumber"]),
                str(batch["lastShotNumber"]),
            ]
        )
        for batch in batches
    ]
    tasks_path.write_text("\n".join(task_lines) + "\n", encoding="utf-8")
    atomic_write_json(
        plan_path,
        {
            "format": PLAN_FORMAT,
            "createdAt": utc_now(),
            "startedEpochSeconds": time.time(),
            "manifestPath": str(manifest_path),
            "manifestSha256": file_sha256(manifest_path),
            "analysisVersion": ANALYSIS_VERSION,
            "batchSize": args.batch_size,
            "shotCount": len(shots),
            "topK": args.top_k,
            "batches": batches,
        },
    )
    progress("analysis-plan-complete", shots=len(shots), batches=len(batches))
    print(
        "SUMMARY|stage=analysis-plan"
        f"|shots={len(shots)}"
        f"|batches={len(batches)}"
        f"|tasksOutput={tasks_path}"
        f"|planOutput={plan_path}",
        flush=True,
    )
    return 0


def manifest_source_without_index(source: dict[str, Any]) -> dict[str, Any]:
    result = copy.deepcopy(source)
    result.pop("index", None)
    return result


def validate_part(
    part: dict[str, Any],
    document: dict[str, Any],
    batch: dict[str, Any],
    schema: dict[str, Any],
    canonical_sources: list[dict[str, Any]] | None,
    expected_analysis_version: str,
) -> list[dict[str, Any]]:
    validate_contract(part, schema)
    if part.get("documentKind") != "ANALYSIS_RESULT":
        raise AutoConformFailure(
            "analysis-finalize",
            "INVALID_BATCH_DOCUMENT",
            "analysis batch is not an ANALYSIS_RESULT document",
            {"batchIndex": batch["batchIndex"]},
        )
    if part.get("analysisVersion") != expected_analysis_version:
        raise AutoConformFailure(
            "analysis-finalize",
            "ANALYSIS_VERSION_MISMATCH",
            "analysis batch version does not match the planned matcher version",
            {"batchIndex": batch["batchIndex"]},
        )
    if part.get("reference") != document.get("reference"):
        raise AutoConformFailure(
            "analysis-finalize",
            "REFERENCE_MISMATCH",
            "analysis batch reference no longer matches the manifest",
            {"batchIndex": batch["batchIndex"]},
        )
    expected_base_sources = [manifest_source_without_index(source) for source in document["sources"]]
    actual_base_sources = [manifest_source_without_index(source) for source in part["sources"]]
    if actual_base_sources != expected_base_sources:
        raise AutoConformFailure(
            "analysis-finalize",
            "SOURCE_MISMATCH",
            "analysis batch source identities no longer match the manifest",
            {"batchIndex": batch["batchIndex"]},
        )
    if canonical_sources is not None and part["sources"] != canonical_sources:
        raise AutoConformFailure(
            "analysis-finalize",
            "SOURCE_INDEX_MISMATCH",
            "analysis batches disagree about source-index metadata",
            {"batchIndex": batch["batchIndex"]},
        )
    actual_ids = [str(shot["shotId"]) for shot in part["shots"]]
    if actual_ids != batch["shotIds"]:
        raise AutoConformFailure(
            "analysis-finalize",
            "SHOT_BATCH_MISMATCH",
            "analysis batch does not contain exactly its planned shot range",
            {"batchIndex": batch["batchIndex"]},
        )
    run = part.get("run") or {}
    if run.get("status") != "COMPLETED" or int(run.get("completedShots", -1)) != len(actual_ids):
        raise AutoConformFailure(
            "analysis-finalize",
            "INCOMPLETE_BATCH",
            "analysis batch does not have a complete terminal run record",
            {"batchIndex": batch["batchIndex"]},
        )
    diagnostic_shot_ids = [
        str(item["shotId"])
        for item in (run.get("analysisDiagnostics") or {}).get("shots") or []
    ]
    if diagnostic_shot_ids != actual_ids:
        raise AutoConformFailure(
            "analysis-finalize",
            "DIAGNOSTIC_SHOT_BATCH_MISMATCH",
            "analysis diagnostics do not cover exactly the ordered batch shot range",
            {"batchIndex": batch["batchIndex"]},
        )
    return part["shots"]


def finalize_analysis(args: argparse.Namespace) -> int:
    manifest_path = Path(args.manifest).resolve()
    plan_path = Path(args.plan_output).resolve()
    parts_dir = Path(args.parts_dir).resolve()
    output_path = Path(args.output).resolve()
    placement_path = Path(args.placement_table_output).resolve()
    alternatives_path = Path(args.alternative_table_output).resolve()
    schema = load_schema(Path(args.schema).resolve())
    document = json.loads(manifest_path.read_text(encoding="utf-8"))
    validate_contract(document, schema)
    validate_unique_shot_ids(document["shots"])
    plan = json.loads(plan_path.read_text(encoding="utf-8"))
    if plan.get("format") != PLAN_FORMAT:
        raise AutoConformFailure(
            "analysis-finalize",
            "INVALID_PLAN",
            "analysis plan format is invalid",
        )
    if plan.get("manifestSha256") != file_sha256(manifest_path):
        raise AutoConformFailure(
            "analysis-finalize",
            "MANIFEST_CHANGED",
            "analysis manifest changed after batch planning",
        )
    expected_batches = batch_plan(document["shots"], int(plan["batchSize"]))
    if plan.get("batches") != expected_batches or int(plan.get("shotCount", -1)) != len(document["shots"]):
        raise AutoConformFailure(
            "analysis-finalize",
            "PLAN_MISMATCH",
            "analysis plan no longer matches the manifest shot list",
        )

    completed_shots: list[dict[str, Any]] = []
    canonical_sources: list[dict[str, Any]] | None = None
    timing_keys = (
        "indexingSeconds",
        "retrievalSeconds",
        "verificationSeconds",
        "refinementSeconds",
        "totalSeconds",
    )
    timing = {key: 0.0 for key in timing_keys}
    analysis_diagnostics = new_analysis_diagnostics()
    peak_memory_bytes = 0
    temporary_disk_bytes = 0
    for batch in expected_batches:
        part_path = parts_dir / str(batch["fileName"])
        if not part_path.is_file():
            raise AutoConformFailure(
                "analysis-finalize",
                "MISSING_BATCH",
                "a planned analysis batch result is missing",
                {"batchIndex": batch["batchIndex"], "path": str(part_path)},
            )
        part = json.loads(part_path.read_text(encoding="utf-8"))
        completed_shots.extend(
            validate_part(
                part,
                document,
                batch,
                schema,
                canonical_sources,
                str(plan["analysisVersion"]),
            )
        )
        if canonical_sources is None:
            canonical_sources = copy.deepcopy(part["sources"])
        run = part["run"]
        for key in timing_keys:
            timing[key] += float(run["timing"][key])
        merge_analysis_diagnostics(
            analysis_diagnostics,
            run["analysisDiagnostics"],
        )
        peak_memory_bytes = max(peak_memory_bytes, int(run["peakMemoryBytes"]))
        temporary_disk_bytes += int(run["temporaryDiskBytes"])

    # Batch process totals omit the orchestration time between Python calls. For an
    # editor-facing processing time, totalSeconds represents wall-clock time from
    # planning through the complete, validated analysis result.
    timing["totalSeconds"] = max(
        timing["totalSeconds"],
        0.0,
        time.time() - float(plan["startedEpochSeconds"]),
    )

    expected_ids = [str(shot["shotId"]) for shot in document["shots"]]
    actual_ids = [str(shot["shotId"]) for shot in completed_shots]
    if actual_ids != expected_ids:
        raise AutoConformFailure(
            "analysis-finalize",
            "INCOMPLETE_ANALYSIS",
            "combined analysis batches do not cover the complete ordered shot list",
        )

    result = copy.deepcopy(document)
    result["documentKind"] = "ANALYSIS_RESULT"
    result["analysisVersion"] = str(plan.get("analysisVersion") or ANALYSIS_VERSION)
    result["sources"] = canonical_sources or []
    result["shots"] = completed_shots
    result["run"] = {
        "status": "COMPLETED",
        "timing": timing,
        "peakMemoryBytes": peak_memory_bytes,
        "temporaryDiskBytes": temporary_disk_bytes,
        "analysisDiagnostics": analysis_diagnostics,
        "completedShots": len(completed_shots),
    }
    result["benchmarkSummary"] = benchmark_summary(completed_shots, int(plan["topK"]))
    validate_contract(result, schema)
    atomic_write_json(output_path, result)
    write_placement_table(placement_path, completed_shots, result["sources"])
    write_alternative_table(alternatives_path, completed_shots, result["sources"])
    progress(
        "analysis-finalize-complete",
        shots=len(completed_shots),
        batches=len(expected_batches),
        output=output_path,
    )
    print(
        "SUMMARY|stage=analysis-finalize"
        f"|shots={len(completed_shots)}"
        f"|batches={len(expected_batches)}"
        f"|output={output_path}",
        flush=True,
    )
    return 0


def structured_error_value(value: object) -> str:
    return str(value).replace("|", "/").replace("\r", " ").replace("\n", " ")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--operation", choices=["PLAN", "FINALIZE"], required=True)
    parser.add_argument("--manifest", required=True)
    parser.add_argument("--schema", required=True)
    parser.add_argument("--plan-output", required=True)
    parser.add_argument("--tasks-output", required=True)
    parser.add_argument("--parts-dir", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--placement-table-output", required=True)
    parser.add_argument("--alternative-table-output", required=True)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--top-k", type=int, default=5)
    args = parser.parse_args()
    if args.batch_size <= 0:
        raise ValueError("--batch-size must be greater than zero")
    if args.top_k <= 0:
        raise ValueError("--top-k must be greater than zero")
    return plan_analysis(args) if args.operation == "PLAN" else finalize_analysis(args)


def cli() -> int:
    try:
        return main()
    except AutoConformFailure as exc:
        failure = exc
    except (OSError, ValueError, TypeError, KeyError, json.JSONDecodeError) as exc:
        failure = AutoConformFailure(
            "analysis-jobs",
            "ANALYSIS_JOBS_FAILED",
            str(exc) or exc.__class__.__name__,
            {"exceptionType": exc.__class__.__name__},
        )
    except Exception as exc:
        failure = AutoConformFailure(
            "analysis-jobs",
            "UNEXPECTED_FAILURE",
            str(exc) or exc.__class__.__name__,
            {"exceptionType": exc.__class__.__name__},
        )
    fields = [
        "ERROR",
        f"stage={structured_error_value(failure.stage)}",
        f"code={structured_error_value(failure.code)}",
        f"message={structured_error_value(failure)}",
    ]
    for key, value in failure.details.items():
        fields.append(f"{key}={structured_error_value(value)}")
    print("|".join(fields), file=sys.stderr, flush=True)
    return 2


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