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

import argparse
import csv
import json
import sys
from pathlib import Path


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

from alignment import analyze_session  # noqa: E402
from contracts import (  # noqa: E402
    validate_analysis,
    validate_markers_tsv,
    validate_placements_tsv,
    validate_source_markers_tsv,
)
from session import plan_session_order, validate_session_order  # noqa: E402


PLACEMENT_COLUMNS = (
    "groupId",
    "groupOrder",
    "occurrenceId",
    "recordingOrder",
    "videoTrackIndex",
    "audioTrackIndex",
    "sequenceStartSeconds",
    "sourceInSeconds",
    "sourceOutSeconds",
    "speechInSeconds",
    "speechOutSeconds",
    "sourceId",
    "sourceOrder",
    "matchClass",
    "similarity",
    "referenceCoverage",
    "occurrenceCoverage",
    "isPartial",
)

MARKER_COLUMNS = (
    "diagnosticId",
    "severity",
    "reasonCode",
    "markerName",
    "markerComment",
    "sequenceStartSeconds",
    "durationSeconds",
    "groupId",
    "referenceUnitId",
)

SOURCE_MARKER_COLUMNS = (
    "sourceMarkerId",
    "sourceId",
    "occurrenceId",
    "groupId",
    "referenceUnitId",
    "status",
    "reasonCode",
    "markerName",
    "markerComment",
    "sourceStartSeconds",
    "durationSeconds",
    "colorIndex",
)

SESSION_ORDER_COLUMNS = (
    "sourceId",
    "sourceOrder",
    "selectionIndex",
    "captureDateDisplay",
    "orderMode",
    "warningCodes",
)

SESSION_INPUT_COLUMNS = (
    "sourceId",
    "selectionIndex",
    "displayName",
    "captureDateRaw",
    "durationSeconds",
    "transcriptPath",
)

SUMMARY_COLUMNS = (
    "mode",
    "orderMode",
    "groupCount",
    "placementCount",
    "maxTrackCount",
    "sequenceMarkerCount",
    "sourceMarkerCount",
)


def _seconds(value_ms: int) -> str:
    return f"{value_ms / 1000.0:.3f}"


def write_placements(analysis: dict, path: Path) -> None:
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=PLACEMENT_COLUMNS, dialect="excel-tab")
        writer.writeheader()
        for group in analysis["groups"]:
            for occurrence in sorted(
                group["occurrences"],
                key=lambda item: (
                    item["videoTrackIndex"],
                    item["audioTrackIndex"],
                    item["recordingOrder"],
                ),
            ):
                writer.writerow(
                    {
                        "groupId": group["groupId"],
                        "groupOrder": group["groupOrder"],
                        "occurrenceId": occurrence["occurrenceId"],
                        "recordingOrder": occurrence["recordingOrder"],
                        "videoTrackIndex": occurrence["videoTrackIndex"],
                        "audioTrackIndex": occurrence["audioTrackIndex"],
                        "sequenceStartSeconds": _seconds(occurrence["sequenceStartMs"]),
                        "sourceInSeconds": _seconds(occurrence["sourceInMs"]),
                        "sourceOutSeconds": _seconds(occurrence["sourceOutMs"]),
                        "speechInSeconds": _seconds(occurrence["speechInMs"]),
                        "speechOutSeconds": _seconds(occurrence["speechOutMs"]),
                        "sourceId": occurrence["sourceId"],
                        "sourceOrder": occurrence["sourceOrder"],
                        "matchClass": occurrence["matchClass"],
                        "similarity": f"{occurrence['similarity']:.6f}",
                        "referenceCoverage": f"{occurrence['referenceCoverage']:.6f}",
                        "occurrenceCoverage": f"{occurrence['occurrenceCoverage']:.6f}",
                        "isPartial": str(occurrence["isPartial"]).lower(),
                    }
                )


def write_markers(analysis: dict, path: Path) -> None:
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=MARKER_COLUMNS, dialect="excel-tab")
        writer.writeheader()
        for diagnostic in analysis["diagnostics"]:
            writer.writerow(
                {
                    "diagnosticId": diagnostic["diagnosticId"],
                    "severity": diagnostic["severity"],
                    "reasonCode": ",".join(diagnostic["reasonCodes"]),
                    "markerName": diagnostic["markerName"],
                    "markerComment": diagnostic["markerComment"],
                    "sequenceStartSeconds": _seconds(diagnostic["sequenceStartMs"]),
                    "durationSeconds": _seconds(diagnostic["durationMs"]),
                    "groupId": diagnostic["groupId"],
                    "referenceUnitId": diagnostic["referenceUnitId"],
                }
            )


def write_source_markers(analysis: dict, path: Path) -> None:
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(
            handle, fieldnames=SOURCE_MARKER_COLUMNS, dialect="excel-tab"
        )
        writer.writeheader()
        for marker in analysis["sourceMarkers"]:
            writer.writerow(
                {
                    "sourceMarkerId": marker["sourceMarkerId"],
                    "sourceId": marker["sourceId"],
                    "occurrenceId": marker["occurrenceId"],
                    "groupId": marker["groupId"],
                    "referenceUnitId": marker["referenceUnitId"],
                    "status": marker["status"],
                    "reasonCode": ",".join(marker["reasonCodes"]),
                    "markerName": marker["markerName"],
                    "markerComment": marker["markerComment"],
                    "sourceStartSeconds": _seconds(marker["sourceStartMs"]),
                    "durationSeconds": _seconds(marker["durationMs"]),
                    "colorIndex": marker["colorIndex"],
                }
            )


def write_session_order(result: dict, path: Path) -> None:
    warning_codes = ",".join(
        warning["reasonCode"] for warning in result["warnings"]
    )
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(
            handle, fieldnames=SESSION_ORDER_COLUMNS, dialect="excel-tab"
        )
        writer.writeheader()
        for source in result["sources"]:
            writer.writerow(
                {
                    **source,
                    "orderMode": result["orderMode"],
                    "warningCodes": warning_codes,
                }
            )


def write_summary(analysis: dict, path: Path) -> None:
    summary = analysis["summary"]
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(
            handle, fieldnames=SUMMARY_COLUMNS, dialect="excel-tab"
        )
        writer.writeheader()
        writer.writerow(
            {
                "mode": analysis["mode"],
                "orderMode": analysis["orderMode"],
                "groupCount": summary["groupCount"],
                "placementCount": summary["occurrenceCount"],
                "maxTrackCount": summary["maxTrackCount"],
                "sequenceMarkerCount": summary["diagnosticCount"],
                "sourceMarkerCount": summary["sourceMarkerCount"],
            }
        )


def load_session_manifest(path: Path, reference_file: Path | None = None) -> dict:
    if path.suffix.casefold() == ".json":
        with path.open("r", encoding="utf-8") as handle:
            manifest = json.load(handle)
    else:
        with path.open("r", encoding="utf-8", newline="") as handle:
            reader = csv.DictReader(handle, dialect="excel-tab")
            if tuple(reader.fieldnames or ()) != SESSION_INPUT_COLUMNS:
                raise ValueError(
                    "Session input table header does not match the frozen contract."
                )
            rows = list(reader)
        sources = []
        for row_index, row in enumerate(rows, start=2):
            try:
                selection_index = int(row["selectionIndex"])
            except ValueError as error:
                raise ValueError(
                    f"Session input row {row_index} selectionIndex must be an integer."
                ) from error
            sources.append(
                {
                    "sourceId": row["sourceId"],
                    "selectionIndex": selection_index,
                    "displayName": row["displayName"],
                    "captureDateRaw": row["captureDateRaw"],
                    "durationSeconds": row["durationSeconds"],
                    "transcriptPath": row["transcriptPath"],
                }
            )
        manifest = {
            "contractVersion": 1,
            "referenceText": "",
            "sources": sources,
        }
    if reference_file is not None:
        manifest["referenceText"] = reference_file.read_text(encoding="utf-8")
    return manifest


def main() -> int:
    parser = argparse.ArgumentParser(description="Offline repeated-take alignment")
    subparsers = parser.add_subparsers(dest="command", required=True)
    plan_parser = subparsers.add_parser("plan-session")
    plan_parser.add_argument("--srt-toolkit-dir", type=Path, help=argparse.SUPPRESS)
    plan_parser.add_argument("--lexical-alignment-dir", type=Path, help=argparse.SUPPRESS)
    plan_parser.add_argument("--session-manifest", required=True, type=Path)
    plan_parser.add_argument("--output", required=True, type=Path)
    plan_parser.add_argument("--output-tsv", type=Path)
    analyze_parser = subparsers.add_parser("analyze")
    analyze_parser.add_argument("--srt-toolkit-dir", type=Path, help=argparse.SUPPRESS)
    analyze_parser.add_argument("--lexical-alignment-dir", type=Path, help=argparse.SUPPRESS)
    analyze_parser.add_argument("--session-manifest", required=True, type=Path)
    analyze_parser.add_argument("--reference-file", type=Path)
    analyze_parser.add_argument("--output-dir", required=True, type=Path)
    analyze_parser.add_argument("--handle-before", type=float, default=0.2)
    analyze_parser.add_argument("--handle-after", type=float, default=0.3)
    analyze_parser.add_argument("--group-gap", type=float, default=0.75)
    validate_parser = subparsers.add_parser("validate-plan")
    validate_parser.add_argument("--srt-toolkit-dir", type=Path, help=argparse.SUPPRESS)
    validate_parser.add_argument("--lexical-alignment-dir", type=Path, help=argparse.SUPPRESS)
    validate_parser.add_argument("--session-manifest", required=True, type=Path)
    validate_parser.add_argument("--analysis", required=True, type=Path)
    validate_parser.add_argument("--placements", required=True, type=Path)
    validate_parser.add_argument("--markers", required=True, type=Path)
    validate_parser.add_argument("--source-markers", required=True, type=Path)
    args = parser.parse_args()
    manifest = load_session_manifest(
        args.session_manifest,
        args.reference_file if args.command == "analyze" else None,
    )
    if args.command == "plan-session":
        result = plan_session_order(manifest)
        validate_session_order(
            result, [source["sourceId"] for source in manifest["sources"]]
        )
        args.output.parent.mkdir(parents=True, exist_ok=True)
        with args.output.open("w", encoding="utf-8") as handle:
            json.dump(result, handle, ensure_ascii=False, indent=2, sort_keys=True)
            handle.write("\n")
        if args.output_tsv:
            args.output_tsv.parent.mkdir(parents=True, exist_ok=True)
            write_session_order(result, args.output_tsv)
        print(
            f"ORDER|mode={result['orderMode']}|sources={len(result['sources'])}"
            f"|warnings={len(result['warnings'])}"
        )
    elif args.command == "analyze":
        manifest["handleBeforeSeconds"] = args.handle_before
        manifest["handleAfterSeconds"] = args.handle_after
        manifest["groupGapSeconds"] = args.group_gap
        analysis = analyze_session(manifest)
        issued_source_ids = [source["sourceId"] for source in manifest["sources"]]
        if set(analysis["sourceOrder"]) != set(issued_source_ids):
            raise ValueError("Analysis sourceOrder is not the issued source permutation.")
        validate_analysis(analysis, analysis["sourceOrder"])
        args.output_dir.mkdir(parents=True, exist_ok=True)
        with (args.output_dir / "analysis.json").open("w", encoding="utf-8") as handle:
            json.dump(analysis, handle, ensure_ascii=False, indent=2, sort_keys=True)
            handle.write("\n")
        write_placements(analysis, args.output_dir / "placements.tsv")
        write_markers(analysis, args.output_dir / "markers.tsv")
        write_source_markers(analysis, args.output_dir / "source-markers.tsv")
        write_summary(analysis, args.output_dir / "summary.tsv")
    else:
        with args.analysis.open("r", encoding="utf-8") as handle:
            analysis = json.load(handle)
        issued_source_ids = [source["sourceId"] for source in manifest["sources"]]
        if set(analysis["sourceOrder"]) != set(issued_source_ids):
            raise ValueError("Analysis sourceOrder is not the issued source permutation.")
        validate_analysis(analysis, analysis["sourceOrder"])
        row_count = validate_placements_tsv(args.placements, PLACEMENT_COLUMNS, analysis)
        marker_count = validate_markers_tsv(args.markers, MARKER_COLUMNS, analysis)
        source_marker_count = validate_source_markers_tsv(
            args.source_markers, SOURCE_MARKER_COLUMNS, analysis
        )
        print(
            f"VALID|groups={len(analysis['groups'])}|placements={row_count}"
            f"|markers={marker_count}|sourceMarkers={source_marker_count}"
        )
    return 0


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