#!/usr/bin/env python3
"""
Run test vectors for Schnorr signature half-aggregation.

Usage:
    python run_test_vectors.py

Expects CSV files generated by gen_test_vectors.py in vectors/ directory.
"""

import csv
import os
from pathlib import Path
import sys
from typing import List, Tuple

from halfagg import Aggregate, IncAggregate, VerifyAggregate


ROOT = Path(__file__).resolve().parent
VECTORS_DIR = ROOT / "vectors"
AGGREGATE_VECTORS = VECTORS_DIR / 'test_vectors_aggregate.csv'
INCAGGREGATE_VECTORS = VECTORS_DIR / 'test_vectors_incaggregate.csv'
VERIFY_VECTORS = VECTORS_DIR / 'test_vectors_verify.csv'


def parse_hex_list(hex_str: str) -> List[bytes]:
    if not hex_str:
        return []
    return [bytes.fromhex(h) for h in hex_str.split(';')]


def parse_pms(pks_str: str, msgs_str: str, sigs_str: str) -> List[Tuple[bytes, bytes, bytes]]:
    pks = parse_hex_list(pks_str)
    msgs = parse_hex_list(msgs_str)
    sigs = parse_hex_list(sigs_str)
    if not pks:
        return []
    return list(zip(pks, msgs, sigs))


def parse_pm(pks_str: str, msgs_str: str) -> List[Tuple[bytes, bytes]]:
    pks = parse_hex_list(pks_str)
    msgs = parse_hex_list(msgs_str)
    if not pks:
        return []
    return list(zip(pks, msgs))


def run_aggregate_tests() -> List[str]:
    """Run Aggregate test vectors. Returns list of failed test descriptions."""
    failures = []

    with open(AGGREGATE_VECTORS, newline='', encoding='utf-8') as f:
        reader = csv.reader(f)
        next(reader)

        for row in reader:
            index, pks_str, msgs_str, sigs_str, expected_result, expected_aggsig, comment = row
            pms = parse_pms(pks_str, msgs_str, sigs_str)

            if expected_result == "FALSE":
                try:
                    Aggregate(pms)
                    failures.append(f"Aggregate #{index}: {comment}")
                except ValueError:
                    pass
            else:
                try:
                    actual = Aggregate(pms)
                    if actual != bytes.fromhex(expected_aggsig):
                        failures.append(f"Aggregate #{index}: {comment}")
                except ValueError:
                    failures.append(f"Aggregate #{index}: {comment}")

    return failures


def run_incaggregate_tests() -> List[str]:
    """Run IncAggregate test vectors. Returns list of failed test descriptions."""
    failures = []

    with open(INCAGGREGATE_VECTORS, newline='', encoding='utf-8') as f:
        reader = csv.reader(f)
        next(reader)

        for row in reader:
            (index, aggsig_str, pm_pks_str, pm_msgs_str,
             pms_pks_str, pms_msgs_str, pms_sigs_str,
             expected_result, expected_aggsig, comment) = row

            aggsig = bytes.fromhex(aggsig_str) if aggsig_str else bytes()
            pm_aggd = parse_pm(pm_pks_str, pm_msgs_str)
            pms_to_agg = parse_pms(pms_pks_str, pms_msgs_str, pms_sigs_str)

            if expected_result == "FALSE":
                try:
                    IncAggregate(aggsig, pm_aggd, pms_to_agg)
                    failures.append(f"IncAggregate #{index}: {comment}")
                except ValueError:
                    pass
            else:
                try:
                    actual = IncAggregate(aggsig, pm_aggd, pms_to_agg)
                    if actual != bytes.fromhex(expected_aggsig):
                        failures.append(f"IncAggregate #{index}: {comment}")
                except ValueError:
                    failures.append(f"IncAggregate #{index}: {comment}")

    return failures


def run_verify_tests() -> List[str]:
    """Run VerifyAggregate test vectors. Returns list of failed test descriptions."""
    failures = []

    with open(VERIFY_VECTORS, newline='', encoding='utf-8') as f:
        reader = csv.reader(f)
        next(reader)

        for row in reader:
            index, aggsig_str, pks_str, msgs_str, expected_result, comment = row

            aggsig = bytes.fromhex(aggsig_str)
            pm = parse_pm(pks_str, msgs_str)

            try:
                result = VerifyAggregate(aggsig, pm)
                expected = (expected_result == "TRUE")
                if result != expected:
                    failures.append(f"VerifyAggregate #{index}: {comment}")
            except ValueError:
                if expected_result == "TRUE":
                    failures.append(f"VerifyAggregate #{index}: {comment}")

    return failures


def main():
    for filename in [AGGREGATE_VECTORS, INCAGGREGATE_VECTORS, VERIFY_VECTORS]:
        if not os.path.exists(filename):
            print(f"ERROR: Missing {filename}. Run gen_test_vectors.py first.")
            sys.exit(1)

    failures = []
    failures.extend(run_aggregate_tests())
    failures.extend(run_incaggregate_tests())
    failures.extend(run_verify_tests())

    if failures:
        print("FAILED:")
        for f in failures:
            print(f"  {f}")
        sys.exit(1)
    else:
        print("All tests have passed.")
        sys.exit(0)


if __name__ == "__main__":
    main()