"""
SPU (Server Performance Unit) parsing utilities.

Utilities for extracting information from SPU identifiers and specifications.
"""

import re
from typing import Optional, Tuple


def extract_machine_generation(spu: str) -> Optional[str]:
    """Extract machine generation from SPU string.

    NOTE: The naming convention is counterintuitive:
    - SPU with "-V2" suffix → Actually Graviton4 (returns "G4")
    - SPU without suffix → Actually Graviton2 (returns "G2")

    Args:
        spu: SPU specification string (e.g., "32C128G-V2", "64C256G")

    Returns:
        Machine generation string (e.g., "G2", "G4") or None if unknown

    Examples:
        >>> extract_machine_generation("32C128G-V2")
        'G4'
        >>> extract_machine_generation("64C256G")
        'G2'
        >>> extract_machine_generation("8C32G-V2")
        'G4'
    """
    if not spu:
        return None

    # Explicit mapping based on suffix pattern
    # SPUs with -V2 suffix are actually Graviton4
    if re.search(r"-V2$", spu, re.IGNORECASE):
        return "G4"

    # SPUs without version suffix are Graviton2
    # This includes patterns like: 2C8G, 32C128G, 64C256G, 8C16G
    if re.match(r"^\d+C\d+G$", spu, re.IGNORECASE):
        return "G2"

    # Unknown pattern
    return None


def extract_cores_from_spu(spu: str) -> Optional[int]:
    """Extract number of cores from SPU string.

    Args:
        spu: SPU specification string (e.g., "32C128G", "64C256G-V2")

    Returns:
        Number of cores or None if not found

    Examples:
        >>> extract_cores_from_spu("32C128G")
        32
        >>> extract_cores_from_spu("64C256G-V2")
        64
    """
    if not spu:
        return None

    # Match pattern: digits followed by C
    match = re.search(r"(\d+)C", spu, re.IGNORECASE)
    if match:
        return int(match.group(1))

    return None


def extract_memory_from_spu(spu: str) -> Optional[int]:
    """Extract memory size in GB from SPU string.

    Args:
        spu: SPU specification string (e.g., "32C128G", "64C256G-V2")

    Returns:
        Memory size in GB or None if not found

    Examples:
        >>> extract_memory_from_spu("32C128G")
        128
        >>> extract_memory_from_spu("64C256G-V2")
        256
    """
    if not spu:
        return None

    # Match pattern: digits followed by G (after C)
    match = re.search(r"C(\d+)G", spu, re.IGNORECASE)
    if match:
        return int(match.group(1))

    return None


def calculate_total_cores(spu: str, count: int) -> int:
    """Calculate total number of cores.

    Args:
        spu: SPU specification string
        count: Number of nodes

    Returns:
        Total cores (cores per node × count)

    Examples:
        >>> calculate_total_cores("32C128G", 2)
        64
        >>> calculate_total_cores("64C256G-V2", 3)
        192
    """
    cores_per_node = extract_cores_from_spu(spu)
    if cores_per_node is None:
        return 0
    return cores_per_node * count


def calculate_total_memory(spu: str, count: int) -> int:
    """Calculate total memory in GB.

    Args:
        spu: SPU specification string
        count: Number of nodes

    Returns:
        Total memory in GB (memory per node × count)

    Examples:
        >>> calculate_total_memory("32C128G", 2)
        256
        >>> calculate_total_memory("64C256G-V2", 3)
        768
    """
    memory_per_node = extract_memory_from_spu(spu)
    if memory_per_node is None:
        return 0
    return memory_per_node * count


def parse_spu_spec(spu: str) -> Tuple[Optional[int], Optional[int], Optional[str]]:
    """Parse SPU specification into components.

    Args:
        spu: SPU specification string

    Returns:
        Tuple of (cores, memory_gb, machine_generation)

    Examples:
        >>> parse_spu_spec("32C128G-V2")
        (32, 128, 'G4')
        >>> parse_spu_spec("64C256G")
        (64, 256, 'G2')
    """
    cores = extract_cores_from_spu(spu)
    memory = extract_memory_from_spu(spu)
    generation = extract_machine_generation(spu)

    return (cores, memory, generation)


def format_spu_description(spu: str, count: int) -> str:
    """Format SPU specification into human-readable description.

    Args:
        spu: SPU specification string
        count: Number of nodes

    Returns:
        Human-readable description

    Examples:
        >>> format_spu_description("32C128G-V2", 2)
        '2 × 32C128G (G4)'
        >>> format_spu_description("64C256G", 3)
        '3 × 64C256G (G2)'
    """
    cores, memory, generation = parse_spu_spec(spu)

    if cores and memory:
        base = f"{count} × {cores}C{memory}G"
        if generation:
            return f"{base} ({generation})"
        return base

    return f"{count} × {spu}"


if __name__ == "__main__":
    # Test the functions
    test_cases = [
        "32C128G-V2",
        "16C64G",
        "8C32G-V2",
    ]

    print("Testing SPU parsing utilities:\n")
    for spu in test_cases:
        print(f"SPU: {spu}")
        print(f"  Machine Generation: {extract_machine_generation(spu)}")
        print(f"  Cores: {extract_cores_from_spu(spu)}")
        print(f"  Memory: {extract_memory_from_spu(spu)} GB")
        print(
            f"  Total (×2): {calculate_total_cores(spu, 2)} cores, {calculate_total_memory(spu, 2)} GB"
        )
        print(f"  Description: {format_spu_description(spu, 2)}")
        print()