#!/usr/bin/env python3
"""Benchmark AT-SPI overhead and coverage.

Usage:
    benchmark.py [--app APP] [--json]

Options:
    --app APP    Benchmark specific app only
    --json       Output as JSON
"""

import argparse
import json
import sys
import time
from dataclasses import dataclass, field, asdict

import pyatspi

from common import set_debug, log_debug


@dataclass
class AppCoverage:
    """Coverage data for an application."""
    app_name: str
    window_count: int = 0
    element_count: int = 0
    roles_found: set = field(default_factory=set)
    has_coordinates: bool = False
    has_text_elements: bool = False
    has_actions: bool = False
    enumeration_time_ms: float = 0.0
    errors: list = field(default_factory=list)

    def to_dict(self):
        d = asdict(self)
        d['roles_found'] = sorted(list(self.roles_found))
        return d


def count_elements(accessible, stats, depth=0, max_depth=20):
    """Recursively count elements and gather stats."""
    if depth > max_depth:
        return

    try:
        stats.element_count += 1
        stats.roles_found.add(accessible.getRoleName())

        # Check for coordinates
        try:
            component = accessible.queryComponent()
            if component:
                rect = component.getExtents(pyatspi.DESKTOP_COORDS)
                if rect.width > 0 and rect.height > 0:
                    stats.has_coordinates = True
        except Exception:
            pass

        # Check for text
        try:
            text_iface = accessible.queryText()
            if text_iface and text_iface.characterCount > 0:
                stats.has_text_elements = True
        except Exception:
            pass

        # Check for actions
        try:
            action_iface = accessible.queryAction()
            if action_iface and action_iface.nActions > 0:
                stats.has_actions = True
        except Exception:
            pass

        # Recurse to children
        for i in range(accessible.childCount):
            try:
                child = accessible.getChildAtIndex(i)
                if child:
                    count_elements(child, stats, depth + 1, max_depth)
            except Exception as e:
                log_debug(f"Error accessing child {i}: {e}")

    except Exception as e:
        stats.errors.append(str(e))
        log_debug(f"Error counting element: {e}")


def benchmark_app(app):
    """Benchmark a single application."""
    stats = AppCoverage(app_name=app.name or "(unnamed)")

    # Count windows
    stats.window_count = app.childCount

    # Time enumeration
    start = time.perf_counter()
    for i in range(app.childCount):
        try:
            window = app.getChildAtIndex(i)
            if window:
                count_elements(window, stats)
        except Exception as e:
            stats.errors.append(f"Window {i}: {e}")
    elapsed = time.perf_counter() - start
    stats.enumeration_time_ms = round(elapsed * 1000, 2)

    return stats


def benchmark_find(desktop, role="button", iterations=5):
    """Benchmark element finding."""
    times = []

    for _ in range(iterations):
        start = time.perf_counter()
        count = 0

        for i in range(desktop.childCount):
            app = desktop.getChildAtIndex(i)
            if not app:
                continue
            for j in range(app.childCount):
                window = app.getChildAtIndex(j)
                if not window:
                    continue
                # Simple search
                count += _count_role(window, role, 0, 15)

        elapsed = time.perf_counter() - start
        times.append(elapsed * 1000)

    return {
        "role": role,
        "iterations": iterations,
        "avg_ms": round(sum(times) / len(times), 2),
        "min_ms": round(min(times), 2),
        "max_ms": round(max(times), 2),
    }


def _count_role(accessible, role, depth, max_depth):
    """Count elements matching role."""
    if depth > max_depth:
        return 0

    count = 0
    try:
        if role.lower() in accessible.getRoleName().lower():
            count = 1

        for i in range(accessible.childCount):
            try:
                child = accessible.getChildAtIndex(i)
                if child:
                    count += _count_role(child, role, depth + 1, max_depth)
            except Exception:
                pass
    except Exception:
        pass

    return count


def print_coverage_report(results, find_benchmark):
    """Print human-readable coverage report."""
    print("=" * 60)
    print("AT-SPI Coverage Audit")
    print("=" * 60)

    # Summary table
    print(f"\n{'App':<25} {'Win':>4} {'Elem':>6} {'Time':>8} {'Coord':>6} {'Text':>5} {'Act':>4}")
    print("-" * 60)

    total_elements = 0
    total_time = 0

    for r in sorted(results, key=lambda x: x.element_count, reverse=True):
        coord = "✓" if r.has_coordinates else "-"
        text = "✓" if r.has_text_elements else "-"
        act = "✓" if r.has_actions else "-"
        print(f"{r.app_name[:24]:<25} {r.window_count:>4} {r.element_count:>6} "
              f"{r.enumeration_time_ms:>7.1f}ms {coord:>6} {text:>5} {act:>4}")
        total_elements += r.element_count
        total_time += r.enumeration_time_ms

    print("-" * 60)
    print(f"{'TOTAL':<25} {'':<4} {total_elements:>6} {total_time:>7.1f}ms")

    # Roles summary
    all_roles = set()
    for r in results:
        all_roles.update(r.roles_found)

    print(f"\nUnique roles found: {len(all_roles)}")
    print(f"Roles: {', '.join(sorted(all_roles)[:20])}")
    if len(all_roles) > 20:
        print(f"  ... and {len(all_roles) - 20} more")

    # Find benchmark
    if find_benchmark:
        print(f"\nFind Benchmark (role='{find_benchmark['role']}'):")
        print(f"  Avg: {find_benchmark['avg_ms']:.1f}ms "
              f"(min: {find_benchmark['min_ms']:.1f}ms, max: {find_benchmark['max_ms']:.1f}ms)")

    # Apps with issues
    apps_with_errors = [r for r in results if r.errors]
    if apps_with_errors:
        print(f"\nApps with AT-SPI errors: {len(apps_with_errors)}")
        for r in apps_with_errors[:5]:
            print(f"  {r.app_name}: {len(r.errors)} errors")


def main():
    parser = argparse.ArgumentParser(description="Benchmark AT-SPI coverage")
    parser.add_argument("--app", "-a", help="Benchmark specific app only")
    parser.add_argument("--json", action="store_true", help="Output as JSON")
    parser.add_argument("--debug", action="store_true", help="Show debug messages")
    parser.add_argument("--skip-find", action="store_true", help="Skip find benchmark")
    args = parser.parse_args()

    if args.debug:
        set_debug(True)

    try:
        desktop = pyatspi.Registry.getDesktop(0)
    except Exception as e:
        print(f"Error accessing AT-SPI: {e}", file=sys.stderr)
        sys.exit(1)

    results = []

    for i in range(desktop.childCount):
        app = desktop.getChildAtIndex(i)
        if not app:
            continue

        if args.app and args.app.lower() not in (app.name or "").lower():
            continue

        stats = benchmark_app(app)
        results.append(stats)

    # Run find benchmark
    find_benchmark = None
    if not args.skip_find:
        find_benchmark = benchmark_find(desktop)

    if args.json:
        output = {
            "coverage": [r.to_dict() for r in results],
            "find_benchmark": find_benchmark,
        }
        print(json.dumps(output, indent=2))
    else:
        print_coverage_report(results, find_benchmark)


if __name__ == "__main__":
    main()
