from __future__ import annotations

import argparse
import json
from collections import defaultdict
from datetime import datetime
from decimal import Decimal, InvalidOperation
from pathlib import Path
from typing import Any

DATE_FORMATS = ("%Y-%m-%d", "%m/%d/%Y", "%b %d, %Y", "%B %d, %Y")
NUMERIC_FIELDS = {"subtotal", "tax_amount", "total_amount", "quantity", "unit_price", "line_total"}
SLICE_FIELDS = ("noise_class", "layout_family", "difficulty")


def normalize_whitespace(value: str) -> str:
    return " ".join(value.split())


def normalize_date(value: str) -> str:
    stripped = normalize_whitespace(value)
    for fmt in DATE_FORMATS:
        try:
            return datetime.strptime(stripped, fmt).date().isoformat()
        except ValueError:
            continue
    return stripped


def normalize_numeric(value: Any) -> str:
    try:
        return f"{Decimal(str(value)):.2f}"
    except (InvalidOperation, ValueError):
        return normalize_whitespace(str(value))


def normalize_value(field_name: str, value: Any) -> str:
    if value is None:
        return ""
    if isinstance(value, dict) and "value" in value:
        value = value["value"]
    if field_name in {"invoice_date", "due_date"}:
        return normalize_date(str(value))
    if field_name in NUMERIC_FIELDS:
        return normalize_numeric(value)
    return normalize_whitespace(str(value))


def is_attempted_value(field_name: str, value: Any) -> bool:
    return normalize_value(field_name, value) != ""


def load_labels(labels_dir: Path) -> dict[str, dict[str, Any]]:
    labels = {}
    for path in sorted(labels_dir.glob("*.json")):
        payload = json.loads(path.read_text(encoding="utf-8"))
        labels[payload["document_id"]] = payload
    return labels


def load_predictions_payload(predictions_path: Path) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    payload = json.loads(predictions_path.read_text(encoding="utf-8"))
    if isinstance(payload, list):
        return payload, {}
    if isinstance(payload, dict) and isinstance(payload.get("predictions"), list):
        return payload["predictions"], payload.get("_meta", {})
    raise ValueError("Predictions JSON must be an array or an object with `predictions` and optional `_meta`.")


def new_slice_bucket() -> dict[str, float]:
    return {
        "invoice_count": 0,
        "field_correct": 0,
        "field_total": 0,
        "line_correct": 0,
        "line_total": 0,
        "invoice_success_count": 0,
    }


def finalize_slice_trackers(trackers: dict[str, dict[str, dict[str, float]]]) -> dict[str, dict[str, dict[str, float]]]:
    finalized = {}
    for slice_name, buckets in trackers.items():
        finalized[slice_name] = {}
        for bucket_name, stats in sorted(buckets.items()):
            invoice_count = int(stats["invoice_count"])
            field_total = int(stats["field_total"])
            line_total = int(stats["line_total"])
            success_count = int(stats["invoice_success_count"])
            finalized[slice_name][bucket_name] = {
                "invoice_count": invoice_count,
                "overall_field_accuracy": round(stats["field_correct"] / field_total, 4) if field_total else 0.0,
                "line_item_accuracy": round(stats["line_correct"] / line_total, 4) if line_total else 0.0,
                "invoice_level_success_rate": round(success_count / invoice_count, 4) if invoice_count else 0.0,
            }
        if slice_name == "noise_class":
            finalized[slice_name].setdefault("clean", {"invoice_count": 0, "overall_field_accuracy": 0.0, "line_item_accuracy": 0.0, "invoice_level_success_rate": 0.0})
            finalized[slice_name].setdefault("light_noisy", {"invoice_count": 0, "overall_field_accuracy": 0.0, "line_item_accuracy": 0.0, "invoice_level_success_rate": 0.0})
    return finalized


def summarize_latency(latencies: list[float]) -> dict[str, float] | None:
    if not latencies:
        return None
    ordered = sorted(latencies)
    midpoint = len(ordered) // 2
    median = ordered[midpoint] if len(ordered) % 2 == 1 else (ordered[midpoint - 1] + ordered[midpoint]) / 2
    return {
        "count": len(ordered),
        "mean_runtime_ms": round(sum(ordered) / len(ordered), 4),
        "median_runtime_ms": round(median, 4),
        "min_runtime_ms": round(min(ordered), 4),
        "max_runtime_ms": round(max(ordered), 4),
    }


def summarize_costs(costs: list[float]) -> dict[str, float] | None:
    if not costs:
        return None
    return {
        "count": len(costs),
        "total_cost_usd": round(sum(costs), 4),
        "mean_cost_per_invoice_usd": round(sum(costs) / len(costs), 4),
        "min_cost_usd": round(min(costs), 4),
        "max_cost_usd": round(max(costs), 4),
    }


def coverage_summary(attempted: int, total: int) -> dict[str, float]:
    return {"attempted": attempted, "total": total, "coverage": round(attempted / total, 4) if total else 0.0}


def evaluate_predictions(predictions_path: Path, labels_dir: Path, output_path: Path | None = None) -> dict[str, Any]:
    predictions, meta = load_predictions_payload(predictions_path)
    labels = load_labels(labels_dir)
    runtime_ms_by_document = meta.get("runtime_ms_by_document", {}) if isinstance(meta, dict) else {}
    cost_usd_by_document = meta.get("cost_usd_by_document", {}) if isinstance(meta, dict) else {}
    per_field_counts = defaultdict(lambda: {"correct": 0, "total": 0})
    slice_trackers = {field: defaultdict(new_slice_bucket) for field in SLICE_FIELDS}
    missing_field_count = 0
    extra_field_count = 0
    header_field_attempted = 0
    header_field_total = 0
    line_item_attempted = 0
    line_item_correct = 0
    line_item_total = 0
    invoice_success_count = 0
    document_breakdown = []
    mismatch_examples = []
    latencies = []
    costs = []

    for prediction in predictions:
        label = labels[prediction["document_id"]]
        predicted_headers = prediction.get("header_fields", {})
        expected_headers = label["header_fields"]
        document_correct = 0
        document_total = 0
        document_attempted = 0
        document_missing = 0
        document_extra = len(set(predicted_headers) - set(expected_headers))
        document_line_attempted = 0
        document_line_correct = 0
        document_line_total = 0
        runtime_ms = prediction.get("runtime_ms", runtime_ms_by_document.get(prediction["document_id"]))
        cost_usd = prediction.get("cost_usd", cost_usd_by_document.get(prediction["document_id"]))
        if runtime_ms is not None:
            latencies.append(float(runtime_ms))
        if cost_usd is not None:
            costs.append(float(cost_usd))

        for field, expected_value in expected_headers.items():
            per_field_counts[field]["total"] += 1
            document_total += 1
            header_field_total += 1
            if field not in predicted_headers:
                missing_field_count += 1
                document_missing += 1
                if len(mismatch_examples) < 10:
                    mismatch_examples.append({"document_id": prediction["document_id"], "kind": "missing_field", "field": field, "expected": expected_value["value"] if isinstance(expected_value, dict) else expected_value, "predicted": None})
                continue
            if is_attempted_value(field, predicted_headers[field]):
                document_attempted += 1
                header_field_attempted += 1
            normalized_predicted = normalize_value(field, predicted_headers[field])
            normalized_expected = normalize_value(field, expected_value)
            if normalized_predicted == normalized_expected:
                per_field_counts[field]["correct"] += 1
                document_correct += 1
            elif len(mismatch_examples) < 10:
                mismatch_examples.append({"document_id": prediction["document_id"], "kind": "field_mismatch", "field": field, "expected": expected_value["value"] if isinstance(expected_value, dict) else expected_value, "predicted": predicted_headers[field]})

        extra_field_count += document_extra
        if document_extra and len(mismatch_examples) < 10:
            for extra_field in sorted(set(predicted_headers) - set(expected_headers)):
                if len(mismatch_examples) >= 10:
                    break
                mismatch_examples.append({"document_id": prediction["document_id"], "kind": "extra_field", "field": extra_field, "expected": None, "predicted": predicted_headers[extra_field]})

        predicted_items = prediction.get("line_items", [])
        expected_items = label["line_items"]
        for item_index, (expected_item, predicted_item) in enumerate(zip(expected_items, predicted_items)):
            for key, expected_value in expected_item.items():
                line_item_total += 1
                document_line_total += 1
                if is_attempted_value(key, predicted_item.get(key)):
                    line_item_attempted += 1
                    document_line_attempted += 1
                if normalize_value(key, predicted_item.get(key)) == normalize_value(key, expected_value):
                    line_item_correct += 1
                    document_line_correct += 1
                elif len(mismatch_examples) < 10:
                    mismatch_examples.append({"document_id": prediction["document_id"], "kind": "line_item_mismatch", "field": f"line_items[{item_index}].{key}", "expected": expected_value, "predicted": predicted_item.get(key)})

        if len(predicted_items) != len(expected_items) and len(mismatch_examples) < 10:
            mismatch_examples.append({"document_id": prediction["document_id"], "kind": "line_item_count_mismatch", "field": "line_items", "expected": len(expected_items), "predicted": len(predicted_items)})

        invoice_success = (
            document_correct == document_total
            and document_missing == 0
            and document_extra == 0
            and len(predicted_items) == len(expected_items)
            and document_line_total > 0
            and document_line_correct == document_line_total
        )
        if invoice_success:
            invoice_success_count += 1

        document_result = {
            "document_id": prediction["document_id"],
            "field_accuracy": round(document_correct / document_total, 4) if document_total else 0.0,
            "field_coverage": round(document_attempted / document_total, 4) if document_total else 0.0,
            "missing_field_count": document_missing,
            "extra_field_count": document_extra,
            "line_item_accuracy": round(document_line_correct / document_line_total, 4) if document_line_total else 0.0,
            "line_item_coverage": round(document_line_attempted / document_line_total, 4) if document_line_total else 0.0,
            "invoice_success": invoice_success,
        }
        if runtime_ms is not None:
            document_result["runtime_ms"] = round(float(runtime_ms), 4)
        if cost_usd is not None:
            document_result["cost_usd"] = round(float(cost_usd), 4)
        document_breakdown.append(document_result)

        for slice_field in SLICE_FIELDS:
            bucket = slice_trackers[slice_field][label[slice_field]]
            bucket["invoice_count"] += 1
            bucket["field_correct"] += document_correct
            bucket["field_total"] += document_total
            bucket["line_correct"] += document_line_correct
            bucket["line_total"] += document_line_total
            bucket["invoice_success_count"] += int(invoice_success)

    total_correct = sum(item["correct"] for item in per_field_counts.values())
    total_fields = sum(item["total"] for item in per_field_counts.values()) or 1
    invoice_count = len(predictions)
    report = {
        "overall_field_accuracy": round(total_correct / total_fields, 4),
        "invoice_count": invoice_count,
        "invoice_level_success_count": invoice_success_count,
        "invoice_level_success_rate": round(invoice_success_count / invoice_count, 4) if invoice_count else 0.0,
        "coverage": {
            "header_fields": coverage_summary(header_field_attempted, header_field_total),
            "line_items": coverage_summary(line_item_attempted, line_item_total),
            "overall": coverage_summary(header_field_attempted + line_item_attempted, header_field_total + line_item_total),
        },
        "per_field": {field: round(counts["correct"] / counts["total"], 4) if counts["total"] else 0.0 for field, counts in sorted(per_field_counts.items())},
        "missing_field_count": missing_field_count,
        "extra_field_count": extra_field_count,
        "line_item_accuracy": {"correct": line_item_correct, "total": line_item_total, "accuracy": round(line_item_correct / line_item_total, 4) if line_item_total else 0.0},
        "slice_summaries": finalize_slice_trackers(slice_trackers),
        "document_breakdown": document_breakdown,
        "mismatch_examples": mismatch_examples,
    }
    latency_summary = summarize_latency(latencies)
    if latency_summary is not None:
        report["latency_summary"] = latency_summary
    cost_summary = summarize_costs(costs)
    if cost_summary is not None:
        report["cost_summary"] = cost_summary
    if output_path is not None:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
    return report


def render_eval_report(report: dict) -> str:
    coverage = report["coverage"]
    lines = [
        "# Evaluation Report",
        "",
        "## Overall Metrics",
        f"- Overall field accuracy: {report['overall_field_accuracy']}",
        f"- Invoice count: {report['invoice_count']}",
        f"- Invoice-level success: {report['invoice_level_success_count']} / {report['invoice_count']} ({report['invoice_level_success_rate']})",
        f"- Missing field count: {report['missing_field_count']}",
        f"- Extra field count: {report['extra_field_count']}",
        f"- Line-item accuracy: {report['line_item_accuracy']['accuracy']}",
        "",
        "## Coverage",
        f"- Field coverage: {coverage['header_fields']['attempted']} / {coverage['header_fields']['total']} ({coverage['header_fields']['coverage']})",
        f"- Line-item field coverage: {coverage['line_items']['attempted']} / {coverage['line_items']['total']} ({coverage['line_items']['coverage']})",
        f"- Overall coverage: {coverage['overall']['attempted']} / {coverage['overall']['total']} ({coverage['overall']['coverage']})",
        "",
        "## Invoice-Level Success",
        "- Definition: all expected header fields must match, no extra header fields may appear, predicted line-item count must match expected count, and line-item accuracy for that invoice must be `1.0`.",
        "",
        "## Slice Summaries",
        "",
        "## Clean vs Noisy Summary",
    ]
    for bucket_name, metrics in report["slice_summaries"]["noise_class"].items():
        lines.append(f"- `{bucket_name}`: invoice_count={metrics['invoice_count']}, overall_field_accuracy={metrics['overall_field_accuracy']}, line_item_accuracy={metrics['line_item_accuracy']}, invoice_level_success_rate={metrics['invoice_level_success_rate']}")
    lines.extend(["", "## Layout Summary"])
    for bucket_name, metrics in report["slice_summaries"]["layout_family"].items():
        lines.append(f"- `{bucket_name}`: invoice_count={metrics['invoice_count']}, overall_field_accuracy={metrics['overall_field_accuracy']}, line_item_accuracy={metrics['line_item_accuracy']}, invoice_level_success_rate={metrics['invoice_level_success_rate']}")
    lines.extend(["", "## Difficulty Summary"])
    for bucket_name, metrics in report["slice_summaries"]["difficulty"].items():
        lines.append(f"- `{bucket_name}`: invoice_count={metrics['invoice_count']}, overall_field_accuracy={metrics['overall_field_accuracy']}, line_item_accuracy={metrics['line_item_accuracy']}, invoice_level_success_rate={metrics['invoice_level_success_rate']}")
    if "latency_summary" in report:
        lines.extend(["", "## Latency Summary", f"- Count: {report['latency_summary']['count']}", f"- Mean runtime (ms): {report['latency_summary']['mean_runtime_ms']}", f"- Median runtime (ms): {report['latency_summary']['median_runtime_ms']}", f"- Min runtime (ms): {report['latency_summary']['min_runtime_ms']}", f"- Max runtime (ms): {report['latency_summary']['max_runtime_ms']}"])
    if "cost_summary" in report:
        lines.extend(["", "## Cost Summary", f"- Count: {report['cost_summary']['count']}", f"- Total cost (USD): {report['cost_summary']['total_cost_usd']}", f"- Mean cost per invoice (USD): {report['cost_summary']['mean_cost_per_invoice_usd']}", f"- Min cost (USD): {report['cost_summary']['min_cost_usd']}", f"- Max cost (USD): {report['cost_summary']['max_cost_usd']}"])
    lines.extend(["", "## Per-Field Accuracy"])
    for field_name, accuracy in sorted(report["per_field"].items()):
        lines.append(f"- `{field_name}`: {accuracy}")
    lines.extend(["", "## Document Breakdown"])
    for item in report["document_breakdown"]:
        runtime_suffix = f", runtime_ms={item['runtime_ms']}" if "runtime_ms" in item else ""
        cost_suffix = f", cost_usd={item['cost_usd']}" if "cost_usd" in item else ""
        lines.append(f"- `{item['document_id']}`: field_accuracy={item['field_accuracy']}, field_coverage={item['field_coverage']}, missing={item['missing_field_count']}, extra={item['extra_field_count']}, line_item_accuracy={item['line_item_accuracy']}, line_item_coverage={item['line_item_coverage']}, invoice_success={item['invoice_success']}{runtime_suffix}{cost_suffix}")
    lines.extend(["", "## Mismatch Examples"])
    if report["mismatch_examples"]:
        for example in report["mismatch_examples"][:10]:
            lines.append(f"- `{example['document_id']}` `{example['kind']}` `{example['field']}`: expected={example['expected']} predicted={example['predicted']}")
    else:
        lines.append("- No mismatches found.")
    return "\n".join(lines) + "\n"


def render_eval_report_html(report: dict) -> str:
    def render_slice_table(title: str, slice_summary: dict) -> str:
        rows = "".join(f"<tr><td>{bucket_name}</td><td>{metrics['invoice_count']}</td><td>{metrics['overall_field_accuracy']}</td><td>{metrics['line_item_accuracy']}</td><td>{metrics['invoice_level_success_rate']}</td></tr>" for bucket_name, metrics in slice_summary.items())
        return f"<h2>{title}</h2><table><thead><tr><th>Slice</th><th>Invoice Count</th><th>Field Accuracy</th><th>Line-Item Accuracy</th><th>Invoice Success</th></tr></thead><tbody>{rows}</tbody></table>"

    coverage = report["coverage"]
    per_field_rows = "".join(f"<tr><td>{field_name}</td><td>{accuracy}</td></tr>" for field_name, accuracy in sorted(report["per_field"].items()))
    document_rows = "".join(f"<tr><td>{item['document_id']}</td><td>{item['field_accuracy']}</td><td>{item['missing_field_count']}</td><td>{item['extra_field_count']}</td><td>{item['field_coverage']}</td><td>{item['line_item_accuracy']}</td><td>{item['line_item_coverage']}</td><td>{item['invoice_success']}</td></tr>" for item in report["document_breakdown"])
    mismatch_rows = "".join(f"<li><code>{example['document_id']}</code> <code>{example['kind']}</code> <code>{example['field']}</code>: expected={example['expected']} predicted={example['predicted']}</li>" for example in report["mismatch_examples"][:10]) or "<li>No mismatches found.</li>"
    latency_block = ""
    if "latency_summary" in report:
        latency_block = f"<h2>Latency Summary</h2><p>Count: {report['latency_summary']['count']} | Mean: {report['latency_summary']['mean_runtime_ms']} ms | Median: {report['latency_summary']['median_runtime_ms']} ms | Min: {report['latency_summary']['min_runtime_ms']} ms | Max: {report['latency_summary']['max_runtime_ms']} ms</p>"
    cost_block = ""
    if "cost_summary" in report:
        cost_block = f"<h2>Cost Summary</h2><p>Count: {report['cost_summary']['count']} | Total: ${report['cost_summary']['total_cost_usd']} | Mean per invoice: ${report['cost_summary']['mean_cost_per_invoice_usd']} | Min: ${report['cost_summary']['min_cost_usd']} | Max: ${report['cost_summary']['max_cost_usd']}</p>"
    return f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <title>Evaluation Report</title>
  <style>
    body {{ font-family: Arial, sans-serif; margin: 32px; color: #1f2937; }}
    h1, h2, h3 {{ color: #0f172a; }}
    table {{ border-collapse: collapse; width: 100%; margin: 12px 0 24px; }}
    th, td {{ border: 1px solid #d1d5db; padding: 8px 10px; text-align: left; }}
    th {{ background: #f3f4f6; }}
    .metrics {{ background: #f8fafc; border: 1px solid #e5e7eb; padding: 16px; margin-bottom: 24px; }}
  </style>
</head>
<body>
  <h1>Evaluation Report</h1>
  <div class="metrics">
    <p>Overall Field Accuracy: {report['overall_field_accuracy']}</p>
    <p>Invoice-Level Success: {report['invoice_level_success_count']} / {report['invoice_count']} ({report['invoice_level_success_rate']})</p>
    <p>Missing Field Count: {report['missing_field_count']} | Extra Field Count: {report['extra_field_count']}</p>
    <p>Line-Item Accuracy: {report['line_item_accuracy']['accuracy']}</p>
  </div>
  <h2>Coverage</h2>
  <p>Field coverage: {coverage['header_fields']['attempted']} / {coverage['header_fields']['total']} ({coverage['header_fields']['coverage']})</p>
  <p>Line-item field coverage: {coverage['line_items']['attempted']} / {coverage['line_items']['total']} ({coverage['line_items']['coverage']})</p>
  <p>Overall coverage: {coverage['overall']['attempted']} / {coverage['overall']['total']} ({coverage['overall']['coverage']})</p>
  <h2>Invoice-Level Success</h2>
  <p>All expected header fields must match, no extra header fields may appear, predicted line-item count must match expected count, and line-item accuracy for that invoice must be 1.0.</p>
  <h2>Slice Summaries</h2>
  {render_slice_table('Clean vs Noisy Summary', report['slice_summaries']['noise_class'])}
  {render_slice_table('Layout Summary', report['slice_summaries']['layout_family'])}
  {render_slice_table('Difficulty Summary', report['slice_summaries']['difficulty'])}
  {latency_block}
  {cost_block}
  <h2>Per-Field Accuracy</h2>
  <table><thead><tr><th>Field</th><th>Accuracy</th></tr></thead><tbody>{per_field_rows}</tbody></table>
  <h2>Document Breakdown</h2>
  <table><thead><tr><th>Document</th><th>Field Accuracy</th><th>Missing</th><th>Extra</th><th>Field Coverage</th><th>Line-Item Accuracy</th><th>Line-Item Coverage</th><th>Invoice Success</th></tr></thead><tbody>{document_rows}</tbody></table>
  <h2>Mismatch Examples</h2>
  <ul>{mismatch_rows}</ul>
</body>
</html>
"""


def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate invoice extraction predictions against label JSON files.")
    parser.add_argument("--predictions", required=True, help="Path to predictions JSON")
    parser.add_argument("--labels-dir", help="Path to labels directory")
    parser.add_argument("--pack-path", help="Path to generated pack directory")
    parser.add_argument("--output", help="Optional JSON report output path")
    parser.add_argument("--write-markdown", help="Optional Markdown report output path")
    parser.add_argument("--write-html", help="Optional HTML report output path")
    args = parser.parse_args()

    if args.pack_path:
        pack_path = Path(args.pack_path)
        labels_dir = pack_path / "labels"
    elif args.labels_dir:
        labels_dir = Path(args.labels_dir)
    else:
        raise SystemExit("Either --labels-dir or --pack-path is required.")

    report = evaluate_predictions(Path(args.predictions), labels_dir, Path(args.output) if args.output else None)
    if args.write_markdown:
        Path(args.write_markdown).write_text(render_eval_report(report), encoding="utf-8")
    if args.write_html:
        Path(args.write_html).write_text(render_eval_report_html(report), encoding="utf-8")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()
