"""Run adversarial exploit scanning against all environments.

Usage:
    python scripts/run_exploit_scan.py --benchmark humaneval --limit 20
    python scripts/run_exploit_scan.py --benchmark mbpp --limit 50
    python scripts/run_exploit_scan.py --all
"""

import argparse
import hashlib
import logging
import sys
import time
from pathlib import Path

# Ensure src/ is importable when running as a script.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / 'src'))

from deepgym import DeepGym, Environment
from deepgym.adversarial import LLMAdversarialTester
from deepgym.exploit_db import ExploitDB, ExploitRecord

logger = logging.getLogger(__name__)

BENCHMARKS = ['humaneval', 'mbpp']
ENV_ROOT = Path(__file__).resolve().parent.parent / 'environments'


def scan_benchmark(
    benchmark: str,
    limit: int | None = None,
    *,
    api_key: str | None = None,
    model: str = 'claude-haiku-4-5-20251001',
    max_attacks: int = 5,
) -> None:
    """Scan all environments in a benchmark for exploit vulnerabilities.

    Args:
        benchmark: Name of the benchmark directory (e.g. humaneval, mbpp).
        limit: Maximum number of environments to test. None means all.
        api_key: Anthropic API key for LLM attacks.
        model: Anthropic model ID.
        max_attacks: Number of LLM exploit attempts per environment.
    """
    dg = DeepGym(mode='local')
    tester = LLMAdversarialTester(
        dg=dg, api_key=api_key, model=model, max_attacks=max_attacks
    )
    db = ExploitDB()

    env_dir = ENV_ROOT / benchmark
    if not env_dir.is_dir():
        print(f'Benchmark directory not found: {env_dir}')
        return

    env_dirs = sorted([d for d in env_dir.iterdir() if d.is_dir()])
    if limit:
        env_dirs = env_dirs[:limit]

    print(f'Scanning {len(env_dirs)} {benchmark} environments...')

    for i, d in enumerate(env_dirs):
        verifier_path = d / 'verifier.py'
        task_path = d / 'task.md'

        if not verifier_path.exists():
            continue

        existing = db.get(d.name)
        if existing:
            print(f'  [{i + 1}/{len(env_dirs)}] {d.name}: already tested, skipping')
            continue

        verifier_code = verifier_path.read_text(encoding='utf-8')
        task = task_path.read_text(encoding='utf-8') if task_path.exists() else d.name

        env = Environment(task=task, verifier_code=verifier_code)

        print(f'  [{i + 1}/{len(env_dirs)}] {d.name}...', end=' ', flush=True)
        start = time.time()

        try:
            report = tester.test(env)

            patterns = _classify_patterns(verifier_code, report)

            record = ExploitRecord(
                verifier_id=d.name,
                benchmark=benchmark,
                exploitable=not report.is_robust,
                patterns=patterns,
                num_attacks=report.attacks_run,
                num_exploits=report.exploits_found,
                max_exploit_score=max(
                    (r.score for r in report.results if r.exploited), default=0.0
                ),
                attacks=[
                    {
                        'strategy': r.strategy,
                        'score': r.score,
                        'exploited': r.exploited,
                        'details': r.details,
                    }
                    for r in report.results
                ],
                verifier_hash=hashlib.sha256(verifier_code.encode()).hexdigest()[:16],
            )

            db.save(record)
            elapsed = time.time() - start
            status = 'EXPLOITABLE' if record.exploitable else 'ROBUST'
            print(
                f'{status} ({record.num_exploits}/{record.num_attacks} exploits, '
                f'{elapsed:.1f}s)'
            )

        except KeyboardInterrupt:
            print('\nInterrupted by user.')
            break
        except Exception as exc:
            logger.error('Error scanning %s: %s', d.name, exc)
            print(f'ERROR: {exc}')
            continue

    _print_summary(benchmark, db)
    db.close()


def _classify_patterns(verifier_code: str, report) -> list[str]:
    """Classify exploit patterns based on verifier code analysis.

    Args:
        verifier_code: Source code of the verifier.
        report: AdversarialReport from the tester.

    Returns:
        Deduplicated list of pattern labels found.
    """
    patterns: list[str] = []
    code_lower = verifier_code.lower()

    if 'random' not in code_lower and 'seed' not in code_lower:
        patterns.append('static-inputs')

    assert_count = code_lower.count('assert') + code_lower.count('==')
    if assert_count < 5:
        patterns.append('few-test-cases')

    if (
        'inspect' not in code_lower
        and 'source' not in code_lower
        and 'ast' not in code_lower
    ):
        patterns.append('output-only-check')

    if 'type(' not in code_lower and 'isinstance' not in code_lower:
        patterns.append('no-type-validation')

    if 'import' not in code_lower or 'forbidden' not in code_lower:
        patterns.append('no-import-restriction')

    if report and not report.is_robust:
        for r in report.results:
            if r.exploited and 'minimal' in r.strategy.lower():
                patterns.append('implementation-agnostic')
            if r.exploited and 'hardcoded' in r.strategy.lower():
                patterns.append('predictable-tests')

    return list(set(patterns))


def _print_summary(benchmark: str, db: ExploitDB) -> None:
    """Print scan summary statistics.

    Args:
        benchmark: Benchmark name for the header.
        db: ExploitDB with results.
    """
    stats = db.stats()
    print(f'\n=== {benchmark.upper()} SCAN RESULTS ===')
    print(f'Tested: {stats["total_tested"]}')
    exploitable = stats['total_exploitable']
    rate = stats['exploit_rate']
    print(f'Exploitable: {exploitable} ({rate:.1%})')
    print('Pattern frequency:')
    for pattern, count in stats.get('pattern_counts', {}).items():
        print(f'  {pattern}: {count}')


def main() -> None:
    """Parse arguments and run the exploit scanner."""
    parser = argparse.ArgumentParser(description='Run adversarial exploit scanning.')
    parser.add_argument(
        '--benchmark',
        type=str,
        help='Benchmark to scan (humaneval, mbpp, etc.)',
    )
    parser.add_argument(
        '--limit',
        type=int,
        default=None,
        help='Max environments to scan.',
    )
    parser.add_argument(
        '--all',
        action='store_true',
        help='Scan all benchmarks.',
    )
    parser.add_argument(
        '--api-key',
        type=str,
        default=None,
        help='Anthropic API key (falls back to ANTHROPIC_API_KEY env var).',
    )
    parser.add_argument(
        '--model',
        type=str,
        default='claude-haiku-4-5-20251001',
        help='Anthropic model to use for LLM attacks.',
    )
    parser.add_argument(
        '--max-attacks',
        type=int,
        default=5,
        help='Number of LLM exploit attempts per environment.',
    )
    args = parser.parse_args()

    logging.basicConfig(level=logging.WARNING, format='%(levelname)s: %(message)s')

    if args.all:
        for bm in BENCHMARKS:
            scan_benchmark(
                bm,
                args.limit,
                api_key=args.api_key,
                model=args.model,
                max_attacks=args.max_attacks,
            )
    elif args.benchmark:
        scan_benchmark(
            args.benchmark,
            args.limit,
            api_key=args.api_key,
            model=args.model,
            max_attacks=args.max_attacks,
        )
    else:
        parser.error('Specify --benchmark or --all')


if __name__ == '__main__':
    main()
