#!/usr/bin/env python3
"""
uos-kernel-coverage — Kernel coverage integration script for UniversalisOS.

Orchestrates the coverage workflow:
1. Instrument kernel source with uos_cins.py
2. Build instrumented kernel
3. Run instrumented firmware in QEMU
4. Extract coverage data from UART output
5. Parse coverage with uos_covparse.py
6. Generate reports with uos_covexport.py

Usage:
    python3 tools/uos-boot-test/uos-kernel-coverage.py --arch armv7
    python3 tools/uos-boot-test/uos-kernel-coverage.py --arch all --report coverage/
"""

import argparse
import os
import re
import subprocess
import sys
import time
from pathlib import Path
from typing import Optional

# ── Configuration ─────────────────────────────────────────────────────

REPO_ROOT = Path(__file__).resolve().parent.parent.parent
KERNEL_DIR = REPO_ROOT / "kernel"
TOOLS_DIR = REPO_ROOT / "tools" / "uos-cover"
BUILD_DIR = KERNEL_DIR / "build"

# Source directories to instrument
SOURCE_DIRS = [
    KERNEL_DIR / "src" / "core",
    KERNEL_DIR / "src" / "test",
]

# Architecture configs
ARCH_CONFIGS = {
    "armv7": {
        "platform": "qemu-arm-virt",
        "qemu_binary": "qemu-system-arm",
        "qemu_flags": "-M virt -cpu cortex-a15 -m 512M -nographic",
    },
    "aarch64": {
        "platform": "qemu-aarch64-virt",
        "qemu_binary": "qemu-system-aarch64",
        "qemu_flags": "-M virt,gic-version=3,virtualization=on -cpu cortex-a53 -m 512M -smp 4 -nographic",
    },
    "riscv": {
        "platform": "polarfire",
        "qemu_binary": "qemu-system-riscv64",
        "qemu_flags": "-machine microchip-icicle-kit -smp 5 -m 2G -nographic -bios none",
    },
}


# ── Helper Functions ──────────────────────────────────────────────────

def run_cmd(cmd: list[str], cwd: Optional[Path] = None, timeout: int = 300) -> tuple[int, str, str]:
    """Run a command and return (returncode, stdout, stderr)."""
    try:
        result = subprocess.run(
            cmd,
            cwd=cwd or REPO_ROOT,
            capture_output=True,
            timeout=timeout,
        )
        return result.returncode, result.stdout.decode(errors="replace"), result.stderr.decode(errors="replace")
    except subprocess.TimeoutExpired:
        return -1, "", f"Command timed out after {timeout}s"
    except Exception as e:
        return -1, "", str(e)


def find_cpp_files(directories: list[Path]) -> list[Path]:
    """Find all .cpp files in the given directories."""
    files = []
    for d in directories:
        if d.exists():
            files.extend(d.rglob("*.cpp"))
    return files


# ── Coverage Workflow ─────────────────────────────────────────────────

def instrument_source(arch: str) -> tuple[bool, list[Path]]:
    """Instrument kernel source files for coverage."""
    print(f"[coverage] Instrumenting source files for {arch}...")
    
    cpp_files = find_cpp_files(SOURCE_DIRS)
    if not cpp_files:
        print("[coverage] No .cpp files found to instrument")
        return False, []
    
    instrumented = []
    errors = []
    
    for cpp_file in cpp_files:
        # Create instrumented output path
        rel_path = cpp_file.relative_to(KERNEL_DIR)
        instrumented_path = KERNEL_DIR / "instrumented" / rel_path
        instrumented_path.parent.mkdir(parents=True, exist_ok=True)
        
        # Create uxsc path
        uxsc_path = instrumented_path.with_suffix(".uxsc")
        
        # Run uos_cins.py
        cmd = [
            sys.executable,
            str(TOOLS_DIR / "uos_cins.py"),
            "--input", str(cpp_file),
            "--output", str(instrumented_path),
            "--uxsc", str(uxsc_path),
            "--profile", "COV_STATEMENTS",
        ]
        
        rc, stdout, stderr = run_cmd(cmd)
        if rc == 0:
            instrumented.append(instrumented_path)
            print(f"  ✓ {rel_path}")
        else:
            errors.append((rel_path, stderr))
            print(f"  ✗ {rel_path}: {stderr[:100]}")
    
    print(f"[coverage] Instrumented {len(instrumented)}/{len(cpp_files)} files")
    if errors:
        print(f"[coverage] {len(errors)} errors occurred")
    
    return len(errors) == 0, instrumented


def build_instrumented(arch: str) -> bool:
    """Build the kernel with instrumented source."""
    print(f"[coverage] Building instrumented kernel for {arch}...")
    
    config = ARCH_CONFIGS[arch]
    
    # For now, we'll instrument in-place and rebuild
    # A proper implementation would use a separate build tree
    rc, stdout, stderr = run_cmd(
        ["make", f"ARCH={arch}", f"PLATFORM={config['platform']}", "clean"],
        cwd=KERNEL_DIR,
    )
    
    rc, stdout, stderr = run_cmd(
        ["make", f"ARCH={arch}", f"PLATFORM={config['platform']}"],
        cwd=KERNEL_DIR,
        timeout=300,
    )
    
    if rc != 0:
        print(f"[coverage] Build failed:")
        print(stderr[-500:] if stderr else "No error output")
        return False
    
    print(f"[coverage] Build successful")
    return True


def run_firmware(arch: str, timeout: int = 15) -> tuple[bool, str]:
    """Run instrumented firmware in QEMU and capture UART output."""
    print(f"[coverage] Running instrumented firmware in QEMU...")
    
    config = ARCH_CONFIGS[arch]
    elf_path = BUILD_DIR / arch / config["platform"] / "universalisos.elf"
    
    if not elf_path.exists():
        print(f"[coverage] ELF not found: {elf_path}")
        return False, ""
    
    cmd = [
        config["qemu_binary"],
    ] + config["qemu_flags"].split() + [
        "-kernel", str(elf_path),
    ]
    
    start_time = time.time()
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            timeout=timeout,
            cwd=KERNEL_DIR,
        )
        output = result.stdout.decode(errors="replace") + result.stderr.decode(errors="replace")
        return True, output
    except subprocess.TimeoutExpired:
        return True, ""  # Timeout is expected; we capture partial output
    except Exception as e:
        return False, str(e)


def extract_coverage_data(uart_output: str) -> Optional[str]:
    """Extract UMAP coverage data from UART output."""
    # Look for UMAP/1 header in output
    umap_start = uart_output.find("UMAP/1")
    if umap_start == -1:
        # Coverage data might not be emitted yet (no in-kernel runtime)
        print("[coverage] No UMAP data found in UART output")
        print("[coverage] This is expected if the kernel doesn't emit coverage data yet")
        return None
    
    # Extract until end or next section
    umap_data = uart_output[umap_start:]
    # Find end marker or take everything
    end_markers = ["=== ", "UniversalisOS", "\r\n\r\n"]
    end_pos = len(umap_data)
    for marker in end_markers:
        pos = umap_data.find(marker, 10)  # Skip header
        if pos != -1 and pos < end_pos:
            end_pos = pos
    
    return umap_data[:end_pos]


def parse_coverage(umap_data: str, arch: str) -> bool:
    """Parse coverage data with uos_covparse.py."""
    print(f"[coverage] Parsing coverage data...")
    
    # Write UMAP data to temp file
    umap_file = BUILD_DIR / arch / "coverage.umap"
    umap_file.parent.mkdir(parents=True, exist_ok=True)
    umap_file.write_text(umap_data)
    
    # Create UMDB file path
    umdb_file = BUILD_DIR / arch / "coverage.umdb"
    
    # Run uos_covparse.py
    cmd = [
        sys.executable,
        str(TOOLS_DIR / "uos_covparse.py"),
        "--map", str(umdb_file) if umdb_file.exists() else str(umap_file),
        "--dataset", f"test-{arch}",
    ]
    
    rc, stdout, stderr = run_cmd(cmd)
    if rc == 0:
        print(f"[coverage] Coverage parsed successfully")
        return True
    else:
        print(f"[coverage] Parse failed: {stderr[:200]}")
        return False


def generate_report(arch: str, output_dir: Path) -> bool:
    """Generate coverage report with uos_covexport.py."""
    print(f"[coverage] Generating coverage report...")
    
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Run uos_covexport.py
    cmd = [
        sys.executable,
        str(TOOLS_DIR / "uos_covexport.py"),
        "--input", str(BUILD_DIR / arch / "coverage.umdb"),
        "--format", "html",
        "--output", str(output_dir / f"coverage-{arch}.html"),
    ]
    
    rc, stdout, stderr = run_cmd(cmd)
    if rc == 0:
        print(f"[coverage] Report generated: {output_dir / f'coverage-{arch}.html'}")
        return True
    else:
        print(f"[coverage] Report generation failed: {stderr[:200]}")
        return False


# ── Main ──────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="Kernel coverage integration for UniversalisOS",
    )
    parser.add_argument(
        "--arch",
        choices=["armv7", "aarch64", "riscv", "all"],
        default="armv7",
        help="Architecture to instrument (default: armv7)",
    )
    parser.add_argument(
        "--report",
        type=str,
        default="coverage",
        help="Output directory for coverage reports (default: coverage/)",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=15,
        help="QEMU timeout in seconds (default: 15)",
    )
    parser.add_argument(
        "--skip-build",
        action="store_true",
        help="Skip build step (use existing ELF)",
    )
    parser.add_argument(
        "--skip-instrument",
        action="store_true",
        help="Skip instrumentation step (use existing instrumented source)",
    )
    
    args = parser.parse_args()
    
    # Determine architectures
    if args.arch == "all":
        archs = list(ARCH_CONFIGS.keys())
    else:
        archs = [args.arch]
    
    report_dir = REPO_ROOT / args.report
    
    print(f"UniversalisOS Kernel Coverage")
    print(f"Repository: {REPO_ROOT}")
    print(f"Architectures: {archs}")
    print(f"Report directory: {report_dir}")
    print()
    
    overall_success = True
    
    for arch in archs:
        print(f"\n{'='*60}")
        print(f"Coverage for {arch}")
        print(f"{'='*60}")
        
        # Step 1: Instrument source
        if not args.skip_instrument:
            success, instrumented = instrument_source(arch)
            if not success:
                print(f"[coverage] Warning: Some files failed to instrument")
        
        # Step 2: Build
        if not args.skip_build:
            if not build_instrumented(arch):
                overall_success = False
                continue
        
        # Step 3: Run firmware
        success, uart_output = run_firmware(arch, args.timeout)
        if not success:
            print(f"[coverage] Failed to run firmware")
            overall_success = False
            continue
        
        # Step 4: Extract coverage data
        umap_data = extract_coverage_data(uart_output)
        if umap_data is None:
            print(f"[coverage] No coverage data available (expected for initial integration)")
            print(f"[coverage] Coverage instrumentation requires libuoscov runtime in kernel")
            # Still report success for the boot test
            continue
        
        # Step 5: Parse coverage
        if not parse_coverage(umap_data, arch):
            overall_success = False
        
        # Step 6: Generate report
        if not generate_report(arch, report_dir):
            overall_success = False
    
    print(f"\n{'='*60}")
    if overall_success:
        print("Coverage analysis completed successfully")
    else:
        print("Coverage analysis completed with errors")
    print(f"{'='*60}")
    
    return 0 if overall_success else 1


if __name__ == "__main__":
    sys.exit(main())