from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any


def load_json(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as handle:
        return json.load(handle)


def build_diff(label: str, baseline: float | int, defended: float | int) -> dict[str, float | int]:
    if isinstance(baseline, float) or isinstance(defended, float):
        delta = float(defended) - float(baseline)
    else:
        delta = int(defended) - int(baseline)
    return {
        "metric": label,
        "baseline": baseline,
        "defended": defended,
        "delta": delta,
    }


def build_breakdown_diff(
    label: str,
    baseline_breakdown: dict[str, Any],
    defended_breakdown: dict[str, Any],
) -> list[dict[str, Any]]:
    keys = sorted(set(baseline_breakdown) | set(defended_breakdown))
    rows: list[dict[str, Any]] = []
    for key in keys:
        base_row = baseline_breakdown.get(key, {})
        defended_row = defended_breakdown.get(key, {})
        base_total = int(base_row.get("total_cases", 0))
        defended_total = int(defended_row.get("total_cases", 0))
        rows.append(
            {
                "label": key,
                "baseline_successful_attacks": int(base_row.get("successful_attacks", 0)),
                "defended_successful_attacks": int(defended_row.get("successful_attacks", 0)),
                "successful_attack_delta": int(defended_row.get("successful_attacks", 0))
                - int(base_row.get("successful_attacks", 0)),
                "baseline_total_cases": base_total,
                "defended_total_cases": defended_total,
                "baseline_attack_success_rate": float(base_row.get("attack_success_rate", 0.0)),
                "defended_attack_success_rate": float(
                    defended_row.get("attack_success_rate", 0.0)
                ),
                "attack_success_rate_delta": float(
                    defended_row.get("attack_success_rate", 0.0)
                )
                - float(base_row.get("attack_success_rate", 0.0)),
            }
        )
    return rows


def compare_summaries(baseline_summary: dict[str, Any], defended_summary: dict[str, Any]) -> dict[str, Any]:
    base_metrics = baseline_summary.get("aggregate_metrics", {})
    defended_metrics = defended_summary.get("aggregate_metrics", {})
    mitigation_metrics = defended_summary.get("mitigation_metrics", {})

    metrics = [
        build_diff("attack_success_rate", base_metrics.get("attack_success_rate", 0.0), defended_metrics.get("attack_success_rate", 0.0)),
        build_diff("successful_attacks", base_metrics.get("successful_attacks", 0), defended_metrics.get("successful_attacks", 0)),
        build_diff("prompt_leakage_cases", base_metrics.get("prompt_leakage_cases", 0), defended_metrics.get("prompt_leakage_cases", 0)),
        build_diff("refusals", base_metrics.get("refusals", 0), defended_metrics.get("refusals", 0)),
        build_diff("manual_review_cases", base_metrics.get("manual_review_cases", 0), defended_metrics.get("manual_review_cases", 0)),
        build_diff("benign_utility_rate", base_metrics.get("benign_utility_rate", 0.0), defended_metrics.get("benign_utility_rate", 0.0)),
        build_diff("false_positive_rate", base_metrics.get("false_positive_rate", 0.0), defended_metrics.get("false_positive_rate", 0.0)),
    ]

    return {
        "baseline_run_id": baseline_summary.get("run_id"),
        "defended_run_id": defended_summary.get("run_id"),
        "baseline_experiment": baseline_summary.get("experiment_name"),
        "defended_experiment": defended_summary.get("experiment_name"),
        "mitigation_strategy": mitigation_metrics.get("strategy"),
        "mitigation_counts": mitigation_metrics,
        "metric_diffs": metrics,
        "by_attack_category": build_breakdown_diff(
            "by_attack_category",
            base_metrics.get("by_attack_category", {}),
            defended_metrics.get("by_attack_category", {}),
        ),
        "by_attack_surface": build_breakdown_diff(
            "by_attack_surface",
            base_metrics.get("by_attack_surface", {}),
            defended_metrics.get("by_attack_surface", {}),
        ),
        "by_source": build_breakdown_diff(
            "by_source",
            base_metrics.get("by_source", {}),
            defended_metrics.get("by_source", {}),
        ),
    }


def render_markdown_summary(comparison: dict[str, Any]) -> str:
    lines = [
        "# Phase 7 Comparison Summary",
        "",
        f"- Baseline: `{comparison['baseline_experiment']}` ({comparison['baseline_run_id']})",
        f"- Defended: `{comparison['defended_experiment']}` ({comparison['defended_run_id']})",
        f"- Mitigation: `{comparison['mitigation_strategy']}`",
        "",
        "## Overall Metrics",
        "",
        "| Metric | Baseline | Defended | Delta |",
        "| --- | ---: | ---: | ---: |",
    ]
    for item in comparison["metric_diffs"]:
        lines.append(
            f"| {item['metric']} | {item['baseline']} | {item['defended']} | {item['delta']} |"
        )

    def append_section(title: str, rows: list[dict[str, Any]]) -> None:
        lines.extend(
            [
                "",
                f"## {title}",
                "",
                "| Label | Baseline ASR | Defended ASR | Delta | Baseline Cases | Defended Cases |",
                "| --- | ---: | ---: | ---: | ---: | ---: |",
            ]
        )
        for row in rows:
            lines.append(
                f"| {row['label']} | {row['baseline_attack_success_rate']:.3f} | "
                f"{row['defended_attack_success_rate']:.3f} | "
                f"{row['attack_success_rate_delta']:.3f} | "
                f"{row['baseline_total_cases']} | {row['defended_total_cases']} |"
            )

    append_section("By Attack Category", comparison["by_attack_category"])
    append_section("By Attack Surface", comparison["by_attack_surface"])
    append_section("By Source", comparison["by_source"])
    return "\n".join(lines) + "\n"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Compare baseline and defended experiment summaries")
    parser.add_argument("--baseline", type=Path, required=True, help="Path to baseline summary.json")
    parser.add_argument("--defended", type=Path, required=True, help="Path to defended summary.json")
    parser.add_argument("--output", type=Path, help="Optional output path for comparison JSON")
    parser.add_argument(
        "--markdown-output",
        type=Path,
        help="Optional output path for a report-ready Markdown summary",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    comparison = compare_summaries(load_json(args.baseline), load_json(args.defended))
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        with args.output.open("w", encoding="utf-8") as handle:
            json.dump(comparison, handle, ensure_ascii=False, indent=2)
    if args.markdown_output:
        args.markdown_output.parent.mkdir(parents=True, exist_ok=True)
        args.markdown_output.write_text(
            render_markdown_summary(comparison),
            encoding="utf-8",
        )

    print(f"Baseline: {comparison['baseline_experiment']} ({comparison['baseline_run_id']})")
    print(f"Defended: {comparison['defended_experiment']} ({comparison['defended_run_id']})")
    print(f"Mitigation: {comparison['mitigation_strategy']}")
    for item in comparison["metric_diffs"]:
        print(
            f"{item['metric']}: baseline={item['baseline']} defended={item['defended']} delta={item['delta']}"
        )
    for row in comparison["by_attack_category"]:
        print(
            "category="
            f"{row['label']} baseline_asr={row['baseline_attack_success_rate']}"
            f" defended_asr={row['defended_attack_success_rate']}"
            f" delta={row['attack_success_rate_delta']}"
        )


if __name__ == "__main__":
    main()