#!/usr/bin/env python
"""Run endpoints integration tests module-by-module with parallel execution.

Each test module runs sequentially, but within each module test cases are
distributed across workers via pytest-xdist.

Usage:
    pipenv run python tests/integration/run_endpoints.py
    pipenv run python tests/integration/run_endpoints.py -n 10
    pipenv run python tests/integration/run_endpoints.py test_account
    pipenv run python tests/integration/run_endpoints.py account::TestAccountTimeseries
    pipenv run python tests/integration/run_endpoints.py -n 8 -- -x --tb=short
"""

import argparse
import subprocess
import sys
import time
from pathlib import Path

TESTS_DIR = Path(__file__).parent / "endpoints"

# Per-module concurrency overrides.
#
# A few modules exercise endpoints that run expensive full-catalog
# aggregations under the employee profile (e.g. /top-metrics,
# /top-sound-recordings). At high parallelism these saturate the QA Snowflake
# warehouse and the edge returns 504s ("too many 504 error responses"), which
# surface as flaky CI failures rather than real regressions. Serialize such
# modules so their queries don't contend with each other.
PER_MODULE_CONCURRENCY = {
    "test_top": 1,
}


def discover_all():
    """Return all test modules as (label, pytest_target) pairs."""
    return [(p.stem, str(p)) for p in sorted(TESTS_DIR.glob("test_*.py"))]


def resolve_targets(names):
    """Resolve target names to (label, pytest_target) pairs.

    Accepts:
        test_account                    → all tests in test_account.py
        account                         → all tests in test_account.py
        test_account::TestClass         → only TestClass in test_account.py
        account::TestClass              → only TestClass in test_account.py
    """
    resolved = []
    for name in names:
        if "::" in name:
            module_part, class_part = name.split("::", 1)
        else:
            module_part, class_part = name, None

        if not module_part.startswith("test_"):
            module_part = f"test_{module_part}"
        if not module_part.endswith(".py"):
            module_part = f"{module_part}.py"

        path = TESTS_DIR / module_part
        if not path.exists():
            print(f"Module not found: {path}")
            sys.exit(1)

        if class_part:
            label = f"{path.stem}::{class_part}"
            pytest_target = f"{path}::{class_part}"
        else:
            label = path.stem
            pytest_target = str(path)

        resolved.append((label, pytest_target))
    return resolved


def run_target(pytest_target, concurrency, pytest_extra_args):
    cmd = [
        sys.executable,
        "-m",
        "pytest",
        pytest_target,
        "-n",
        str(concurrency),
        "-v",
        *pytest_extra_args,
    ]
    result = subprocess.run(cmd)
    return result.returncode


def main():
    parser = argparse.ArgumentParser(
        description="Run endpoints tests in parallel per module.",
    )
    parser.add_argument(
        "-n",
        "--concurrency",
        type=int,
        default=5,
        help="parallel workers per module (default: 5)",
    )
    parser.add_argument(
        "targets",
        nargs="*",
        help="modules or classes to run (e.g. account account::TestClass)",
    )

    # Split on "--" to separate our args from pytest args
    argv = sys.argv[1:]
    if "--" in argv:
        split_idx = argv.index("--")
        our_argv, pytest_extra = argv[:split_idx], argv[split_idx + 1 :]
    else:
        our_argv, pytest_extra = argv, []

    args = parser.parse_args(our_argv)

    targets = resolve_targets(args.targets) if args.targets else discover_all()

    print(f"Running {len(targets)} target(s), concurrency={args.concurrency}\n")

    results = {}
    overall_start = time.time()

    for label, pytest_target in targets:
        # Honor per-module concurrency overrides. The label is either the module
        # stem ("test_top") or "test_top::TestClass"; match on the module stem.
        module_stem = label.split("::", 1)[0]
        concurrency = PER_MODULE_CONCURRENCY.get(module_stem, args.concurrency)
        print(f"{'=' * 60}")
        print(f"  {label} (concurrency={concurrency})")
        print(f"{'=' * 60}")
        start = time.time()
        rc = run_target(pytest_target, concurrency, pytest_extra)
        elapsed = time.time() - start
        results[label] = (rc, elapsed)
        print()

    total_elapsed = time.time() - overall_start

    print(f"{'=' * 60}")
    print(f"  SUMMARY")
    print(f"{'=' * 60}")
    for label, (rc, elapsed) in results.items():
        status = "PASSED" if rc == 0 else "FAILED"
        print(f"  {status}  {label} ({elapsed:.1f}s)")
    print(f"{'─' * 60}")
    print(f"  Total: {total_elapsed:.1f}s")

    failed = [n for n, (rc, _) in results.items() if rc != 0]
    if failed:
        print(f"  {len(failed)} target(s) failed: {', '.join(failed)}")
        sys.exit(1)
    else:
        print(f"  All {len(results)} target(s) passed")


if __name__ == "__main__":
    main()