"""JSON parser for rocprof-sys outputs.

Parses JSON files generated by rocprof-sys-run:
- wall_clock-[PID].json - Wall clock timing data
- metadata-[PID].json - Run metadata
- functions-[PID].json - Function call data
"""

import json
from pathlib import Path
from typing import Any, Optional

from wafer_core.lib.rocprofiler.systems.types import SystemMetrics


def parse_wall_clock_json(file_path: Path) -> list[SystemMetrics]:
    """Parse wall_clock-[PID].json file.

    Args:
        file_path: Path to wall_clock JSON file

    Returns:
        List of SystemMetrics with timing data
    """
    with open(file_path, "r") as f:
        data = json.load(f)

    metrics = []

    # rocprof-sys wall_clock format varies, handle common structures
    if isinstance(data, list):
        # Array of timing records
        for record in data:
            if not isinstance(record, dict):
                continue

            metrics.append(
                SystemMetrics(
                    function_name=record.get("name", "unknown"),
                    call_count=record.get("count"),
                    total_time_ns=_to_nanoseconds(record.get("total")),
                    mean_time_ns=_to_nanoseconds(record.get("mean")),
                    min_time_ns=_to_nanoseconds(record.get("min")),
                    max_time_ns=_to_nanoseconds(record.get("max")),
                    stddev_ns=_to_nanoseconds(record.get("stddev")),
                )
            )
    elif isinstance(data, dict):
        # Dictionary with timing data
        for name, timing in data.items():
            if not isinstance(timing, dict):
                continue

            metrics.append(
                SystemMetrics(
                    function_name=name,
                    call_count=timing.get("count"),
                    total_time_ns=_to_nanoseconds(timing.get("total")),
                    mean_time_ns=_to_nanoseconds(timing.get("mean")),
                    min_time_ns=_to_nanoseconds(timing.get("min")),
                    max_time_ns=_to_nanoseconds(timing.get("max")),
                    stddev_ns=_to_nanoseconds(timing.get("stddev")),
                )
            )

    return metrics


def parse_metadata_json(file_path: Path) -> dict[str, Any]:
    """Parse metadata-[PID].json file.

    Args:
        file_path: Path to metadata JSON file

    Returns:
        Dictionary with metadata (PID, hostname, command, etc.)
    """
    import re

    with open(file_path, "r") as f:
        data = json.load(f)

    # rocprof-sys uses nested structure: {"rocprofiler-systems": {"metadata": {...}}}
    # Extract the actual metadata from the nested structure
    metadata_dict = data
    if "rocprofiler-systems" in data:
        if "metadata" in data["rocprofiler-systems"]:
            metadata_dict = data["rocprofiler-systems"]["metadata"]
        else:
            metadata_dict = data["rocprofiler-systems"]

    # Extract info section (rocprof-sys stores system info here)
    info = metadata_dict.get("info", {})
    settings = metadata_dict.get("settings", {})

    # Try to extract PID from filename (metadata-[PID].json)
    pid = None
    pid_match = re.search(r"metadata-(\d+)\.json", file_path.name)
    if pid_match:
        pid = int(pid_match.group(1))

    # Extract relevant metadata
    # Note: rocprof-sys metadata structure:
    # - PID: extracted from filename (metadata-[PID].json)
    # - No explicit hostname field (rocprof-sys doesn't capture it)
    # - No explicit command field (only PWD is available)
    metadata = {
        "pid": pid,
        "hostname": None,  # Not captured by rocprof-sys
        "command": None,  # Not captured by rocprof-sys
        "user": info.get("USER"),
        "working_directory": info.get("PWD"),
        "cpu_model": info.get("CPU_MODEL"),
        "launch_date": info.get("LAUNCH_DATE"),
        "launch_time": info.get("LAUNCH_TIME"),
        "rocm_version": info.get("ROCPROFSYS_ROCM_VERSION"),
        "rocprof_version": info.get("ROCPROFSYS_VERSION"),
        "system_name": info.get("ROCPROFSYS_SYSTEM_NAME"),
        "system_processor": info.get("ROCPROFSYS_SYSTEM_PROCESSOR"),
    }

    # Include the full info and settings dictionaries
    metadata["info"] = info
    metadata["settings"] = settings

    # Include any other top-level keys from metadata_dict
    for key, value in metadata_dict.items():
        if key not in metadata and key not in ["info", "settings"]:
            metadata[key] = value

    return metadata


def parse_functions_json(file_path: Path) -> list[SystemMetrics]:
    """Parse functions-[PID].json file.

    Args:
        file_path: Path to functions JSON file

    Returns:
        List of SystemMetrics with function call data
    """
    with open(file_path, "r") as f:
        data = json.load(f)

    metrics = []

    # Handle different JSON structures
    if isinstance(data, list):
        for record in data:
            if not isinstance(record, dict):
                continue

            # Extract function metrics
            metrics.append(
                SystemMetrics(
                    function_name=record.get("name", record.get("function", "unknown")),
                    call_count=record.get("calls", record.get("count")),
                    total_time_ns=_to_nanoseconds(record.get("total_time")),
                    mean_time_ns=_to_nanoseconds(record.get("mean_time")),
                    min_time_ns=_to_nanoseconds(record.get("min_time")),
                    max_time_ns=_to_nanoseconds(record.get("max_time")),
                    cpu_time_ns=_to_nanoseconds(record.get("cpu_time")),
                    gpu_time_ns=_to_nanoseconds(record.get("gpu_time")),
                )
            )
    elif isinstance(data, dict):
        # Dictionary-based format
        for name, func_data in data.items():
            if not isinstance(func_data, dict):
                continue

            metrics.append(
                SystemMetrics(
                    function_name=name,
                    call_count=func_data.get("calls", func_data.get("count")),
                    total_time_ns=_to_nanoseconds(func_data.get("total_time")),
                    mean_time_ns=_to_nanoseconds(func_data.get("mean_time")),
                    min_time_ns=_to_nanoseconds(func_data.get("min_time")),
                    max_time_ns=_to_nanoseconds(func_data.get("max_time")),
                    cpu_time_ns=_to_nanoseconds(func_data.get("cpu_time")),
                    gpu_time_ns=_to_nanoseconds(func_data.get("gpu_time")),
                )
            )

    return metrics


def _to_nanoseconds(value: Any) -> Optional[float]:
    """Convert time value to nanoseconds.

    Handles various time units and formats from rocprof-sys.

    Args:
        value: Time value (can be float, int, string with unit)

    Returns:
        Time in nanoseconds, or None if invalid
    """
    if value is None:
        return None

    if isinstance(value, (int, float)):
        # Assume nanoseconds if no unit specified
        return float(value)

    if isinstance(value, str):
        # Parse strings like "1.5 ms", "500 us", etc.
        value = value.strip().lower()
        try:
            # Split number and unit
            parts = value.split()
            if len(parts) == 2:
                number, unit = parts
                number = float(number)

                # Convert to nanoseconds
                if unit in ["ns", "nanoseconds"]:
                    return number
                elif unit in ["us", "microseconds"]:
                    return number * 1_000
                elif unit in ["ms", "milliseconds"]:
                    return number * 1_000_000
                elif unit in ["s", "seconds"]:
                    return number * 1_000_000_000
            else:
                # No unit, assume nanoseconds
                return float(value)
        except ValueError:
            pass

    return None
