#!/usr/bin/env python3
"""
CoreMark-PRO CPU Performance Benchmark Tool
Runs CoreMark-PRO benchmark at different thread counts (1 to max CPU cores)

This version follows the RDK Docker Tools protocol:
- Input: /workspace/input/task.json (optional configuration)
- Output: /workspace/output/result.json
"""

import argparse
import json
import os
import platform
import psutil
import re
import subprocess
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional

# CoreMark-PRO workloads
WORKLOADS = [
    "cjpeg-rose7-preset",
    "core",
    "linear_alg-mid-100x100-sp",
    "loops-all-mid-10k-sp",
    "nnet_test",
    "parser-125k",
    "radix2-big-64k",
    "sha-test",
    "zip-test"
]

# Reference scores for CoreMark-PRO score calculation
REFERENCE_SCORES = {
    "cjpeg-rose7-preset": 40.3438,
    "core": 2855.0,
    "linear_alg-mid-100x100-sp": 38.5624,
    "loops-all-mid-10k-sp": 0.87959,
    "nnet_test": 1.45853,
    "parser-125k": 4.81116,
    "radix2-big-64k": 99.6587,
    "sha-test": 48.5201,
    "zip-test": 21.3618
}

# Scale factors
SCALE_FACTORS = {
    "cjpeg-rose7-preset": 1,
    "core": 10000,
    "linear_alg-mid-100x100-sp": 1,
    "loops-all-mid-10k-sp": 1,
    "nnet_test": 1,
    "parser-125k": 1,
    "radix2-big-64k": 1,
    "sha-test": 1,
    "zip-test": 1
}


def get_cpu_info() -> dict:
    """Get CPU information."""
    info = {
        "cpu_count": psutil.cpu_count(logical=True),
        "cpu_count_physical": psutil.cpu_count(logical=False),
        "cpu_freq_mhz": None,
        "cpu_percent": psutil.cpu_percent(interval=1),
    }

    # Try to get CPU frequency
    freq = psutil.cpu_freq()
    if freq:
        info["cpu_freq_mhz"] = freq.current

    # Try to get CPU model name from /proc/cpuinfo
    try:
        with open("/proc/cpuinfo", "r") as f:
            for line in f:
                if line.startswith("model name") or line.startswith("Model") or line.startswith("Processor"):
                    info["cpu_model"] = line.split(":")[1].strip()
                    break
                if line.startswith("CPU implementer"):
                    info["cpu_implementer"] = line.split(":")[1].strip()
                if line.startswith("CPU architecture"):
                    info["cpu_arch"] = line.split(":")[1].strip()
    except:
        pass

    if "cpu_model" not in info:
        info["cpu_model"] = platform.processor() or "Unknown"

    # Get architecture
    info["architecture"] = platform.machine()

    return info


def get_memory_info() -> dict:
    """Get memory information."""
    mem = psutil.virtual_memory()
    return {
        "total_mb": mem.total // (1024 * 1024),
        "available_mb": mem.available // (1024 * 1024),
        "percent_used": mem.percent,
        "used_mb": mem.used // (1024 * 1024),
        "free_mb": mem.free // (1024 * 1024)
    }


def get_system_info() -> dict:
    """Get system information."""
    return {
        "hostname": platform.node(),
        "platform": platform.platform(),
        "system": platform.system(),
        "release": platform.release(),
        "version": platform.version(),
        "python_version": platform.python_version(),
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    }


def run_workload(workload: str, thread_num: int, build_dir: Path) -> dict:
    """Run a single CoreMark-PRO workload with specified thread count."""
    exe_path = build_dir / f"{workload}.exe"

    if not exe_path.exists():
        return {
            "workload": workload,
            "thread_num": thread_num,
            "error": f"Executable not found: {exe_path}",
            "iter_per_sec": None
        }

    cmd = [str(exe_path), "-v0", f"-c{thread_num}"]

    start_time = time.time()
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
    elapsed_time = time.time() - start_time

    output = result.stdout + result.stderr

    # Parse iterations per second from output
    # CoreMark-PRO output format: "workloads/sec" at the end
    iter_per_sec = None

    # Try to find "workloads/sec" pattern
    m = re.search(r"workloads/sec\s*[:\s=]+\s*([\d.]+)", output, re.IGNORECASE)
    if m:
        iter_per_sec = float(m.group(1))
    else:
        # Alternative pattern - look for any decimal number followed by workloads/sec
        m = re.search(r"([\d.]+)\s+workloads/sec", output, re.IGNORECASE)
        if m:
            iter_per_sec = float(m.group(1))

    # Try to calculate from time and iterations if not found
    if iter_per_sec is None:
        time_match = re.search(r"time\(secs?\)\s*=\s*([\d.]+)", output, re.IGNORECASE)
        iter_match = re.search(r"Total workload.*?=\s*(\d+)", output, re.IGNORECASE)
        if time_match and iter_match:
            iter_per_sec = float(iter_match.group(1)) / float(time_match.group(1))

    return {
        "workload": workload,
        "thread_num": thread_num,
        "iter_per_sec": iter_per_sec,
        "returncode": result.returncode,
        "elapsed_time_sec": elapsed_time,
        "raw_output": output[:500] if iter_per_sec is None else None
    }


def calculate_coremark_pro_score(results: List[dict]) -> Optional[float]:
    """Calculate CoreMark-PRO score from workload results."""
    import math

    ratios = []
    for r in results:
        workload = r["workload"]
        score = r.get("iter_per_sec")
        if score is None or workload not in REFERENCE_SCORES:
            continue

        ref = REFERENCE_SCORES[workload]
        scale = SCALE_FACTORS[workload]

        ratio = (score / ref) * scale
        ratios.append(ratio)

    if not ratios:
        return None

    # Calculate geometric mean
    product = 1.0
    for r in ratios:
        product *= r

    geometric_mean = product ** (1.0 / len(ratios))
    coremark_pro_score = geometric_mean * 1000

    return coremark_pro_score


def print_results_table(all_results: List[dict], thread_counts: List[int]):
    """Print results in a human-readable table format."""
    print("\n" + "=" * 90)
    print("CoreMark-PRO Workload Results")
    print("=" * 90)

    # Group by workload
    workloads = {}
    for r in all_results:
        w = r["workload"]
        if w not in workloads:
            workloads[w] = {}
        workloads[w][r["thread_num"]] = r.get("iter_per_sec")

    # Header
    col_width = 22
    thread_col = 12

    header = f"{'Workload':<{col_width}}"
    for t in thread_counts:
        header += f"{'T=' + str(t):>{thread_col}}"
    header += f"{'Scaling':>{thread_col}}"
    print(header)
    print("-" * 90)

    # Data rows
    for workload in WORKLOADS:
        if workload not in workloads:
            continue

        row_data = workloads[workload]
        single_core = row_data.get(1)

        row = f"{workload:<{col_width}}"
        for t in thread_counts:
            score = row_data.get(t)
            if score is not None:
                row += f"{score:>{thread_col}.2f}"
            else:
                row += f"{'N/A':>{thread_col}}"

        # Calculate scaling (multi-core / single-core)
        max_threads = max(thread_counts)
        multi_core = row_data.get(max_threads)
        if single_core is not None and multi_core is not None and single_core > 0:
            scaling = multi_core / single_core
            row += f"{scaling:>{thread_col}.2f}"
        else:
            row += f"{'N/A':>{thread_col}}"

        print(row)

    print("=" * 90)


def print_summary(all_results: List[dict], thread_counts: List[int], coremark_pro_score: Optional[float]):
    """Print summary of benchmark results."""
    print("\n" + "=" * 90)
    print("CoreMark-PRO Score Summary")
    print("=" * 90)

    # Group by thread count
    by_thread = {}
    for r in all_results:
        t = r["thread_num"]
        if t not in by_thread:
            by_thread[t] = []
        by_thread[t].append(r)

    print(f"{'Threads':<10} {'CoreMark-PRO Score':>20}")
    print("-" * 90)

    for t in thread_counts:
        if t in by_thread:
            score = calculate_coremark_pro_score(by_thread[t])
            if score is not None:
                print(f"{t:<10} {score:>20.2f}")
            else:
                print(f"{t:<10} {'N/A':>20}")

    if coremark_pro_score is not None:
        print("-" * 90)
        print(f"{'Final CoreMark-PRO Score':<30} {coremark_pro_score:>20.2f}")

    print("=" * 90)


def print_system_info(info: dict):
    """Print system information."""
    print("\n" + "=" * 90)
    print("System Information")
    print("=" * 90)

    cpu = info["cpu"]
    mem = info["memory"]
    sys_info = info["system"]

    print(f"Hostname:        {sys_info['hostname']}")
    print(f"Platform:        {sys_info['platform']}")
    print(f"Architecture:    {cpu['architecture']}")
    print(f"CPU Model:       {cpu.get('cpu_model', 'Unknown')}")
    print(f"CPU Count:       {cpu['cpu_count']} logical, {cpu['cpu_count_physical']} physical")
    if cpu.get('cpu_freq_mhz'):
        print(f"CPU Frequency:   {cpu['cpu_freq_mhz']:.0f} MHz")
    print(f"Memory Total:    {mem['total_mb']} MB")
    print(f"Memory Used:     {mem['percent_used']:.1f}%")
    print(f"Timestamp:       {sys_info['timestamp']}")

    print("=" * 90)


def main():
    parser = argparse.ArgumentParser(description="CoreMark-PRO CPU Performance Benchmark")
    parser.add_argument("--input", help="Input JSON file path (optional)")
    parser.add_argument("--output", required=True, help="Output JSON file path")
    parser.add_argument("--build-dir", default="/workspace/coremark-pro/builds/linux64/gcc64/bin",
                        help="Directory containing CoreMark-PRO binaries")
    args = parser.parse_args()

    build_dir = Path(args.build_dir)

    if not build_dir.exists():
        print(f"[ERROR] Build directory not found: {build_dir}", file=sys.stderr)
        sys.exit(1)

    # Load optional configuration from input
    config = {}
    if args.input and Path(args.input).exists():
        try:
            with open(args.input) as f:
                config = json.load(f)
            print(f"[INFO] Loaded configuration from: {args.input}")
        except Exception as e:
            print(f"[WARNING] Failed to load input config: {e}", file=sys.stderr)

    # Get system information
    cpu_info = get_cpu_info()
    memory_info = get_memory_info()
    system_info = get_system_info()

    # Print system info
    info = {
        "cpu": cpu_info,
        "memory": memory_info,
        "system": system_info
    }
    print_system_info(info)

    # Determine thread counts to test (1 to max CPU count)
    max_cpus = config.get("max_threads", cpu_info["cpu_count"])
    thread_counts = list(range(1, max_cpus + 1))

    print(f"\n[Benchmark] Running CoreMark-PRO with thread counts: {thread_counts}")
    print(f"[Benchmark] Total workloads: {len(WORKLOADS)}")

    all_results = []

    # Run benchmarks for each thread count
    for thread_num in thread_counts:
        print(f"\n--- Running with {thread_num} thread(s) ---")

        for workload in WORKLOADS:
            print(f"  [{thread_num} threads] {workload} ...", end=" ", flush=True)

            try:
                result = run_workload(workload, thread_num, build_dir)
                all_results.append(result)

                if result.get("iter_per_sec") is not None:
                    print(f"{result['iter_per_sec']:.2f} iter/s")
                else:
                    print(f"FAILED (exit code: {result.get('returncode')})")
                    if result.get("raw_output"):
                        print(f"    Output: {result['raw_output'][:200]}")
            except Exception as e:
                print(f"EXCEPTION: {e}")
                all_results.append({
                    "workload": workload,
                    "thread_num": thread_num,
                    "error": str(e),
                    "iter_per_sec": None
                })

    # Calculate final CoreMark-PRO score for max threads
    by_thread = {}
    for r in all_results:
        t = r["thread_num"]
        if t not in by_thread:
            by_thread[t] = []
        by_thread[t].append(r)

    coremark_pro_score = None
    if max_cpus in by_thread:
        coremark_pro_score = calculate_coremark_pro_score(by_thread[max_cpus])

    # Print results tables
    print_results_table(all_results, thread_counts)
    print_summary(all_results, thread_counts, coremark_pro_score)

    # Prepare output
    output = {
        "system_info": info,
        "benchmark_config": {
            "workloads": WORKLOADS,
            "thread_counts": thread_counts,
            "max_threads": max_cpus
        },
        "results": all_results,
        "coremark_pro_score": coremark_pro_score
    }

    # Save to output file
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with open(output_path, "w") as f:
        json.dump(output, f, indent=2)

    print(f"\n[OK] Results saved to: {args.output}")

    if coremark_pro_score is not None:
        print(f"[OK] CoreMark-PRO Score: {coremark_pro_score:.2f}")
        return 0
    else:
        print("[WARNING] Could not calculate CoreMark-PRO score")
        return 1


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