#!/usr/bin/env python3
"""Job-topology helpers for AI process grouping and straggler detection."""

import re
from statistics import median
from typing import Any, Dict, List, Optional

_RUN_ID_PATTERNS = (
    r"(?:--rdzv-id|--run-id|--run_id|--job-id|--job_id)\s+([^\s]+)",
    r"(?:rdzv_id|run_id|job_id|slurm_job_id|ray_job_id)\s*=\s*([^\s]+)",
)
_RANK_PATTERNS = (
    r"(?:--rank|--node-rank|--node_rank|--local-rank|--local_rank)\s+(-?\d+)",
    r"(?:^|\s)(?:rank|node_rank|local_rank)\s*=\s*(-?\d+)(?:\s|$)",
)
_WORLD_SIZE_PATTERNS = (
    r"(?:--world-size|--world_size|--nnodes)\s+(\d+)",
    r"(?:^|\s)(?:world_size|nnodes)\s*=\s*(\d+)(?:\s|$)",
)
_SCRIPT_PATTERN = re.compile(r"([A-Za-z0-9._/-]+\.py)\b")
_MASTER_PATTERN = re.compile(
    r"(?:--master-addr|--master_addr|--rdzv-endpoint|master_addr)\s*=?\s*([^\s:]+)"
)
_LAUNCHER_HINTS = ("torchrun", "deepspeed", "accelerate launch", "horovodrun", "mpirun")


def _safe_float(value: Any) -> float:
    """Convert arbitrary input to float with safe fallback."""
    try:
        numeric = float(value)
    except (TypeError, ValueError):
        return 0.0
    return max(0.0, min(100.0, numeric))


def _safe_int(value: Any) -> Optional[int]:
    """Convert arbitrary input to int with safe fallback."""
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


def _first_match(patterns: tuple[str, ...], haystack: str) -> Optional[str]:
    """Return first regex capture match from the provided patterns."""
    for pattern in patterns:
        match = re.search(pattern, haystack, re.IGNORECASE)
        if match:
            raw_value = match.group(1).strip("\"' ")
            if raw_value:
                return raw_value
    return None


def _slug(value: str, max_length: int = 40) -> str:
    """Create a stable slug for compact job identifiers."""
    cleaned = re.sub(r"[^a-zA-Z0-9_.-]+", "-", value.strip())
    compact = cleaned.strip("-_.")
    if not compact:
        return "unknown"
    return compact[:max_length]


def _infer_job_id(process: Dict[str, Any]) -> str:
    """Infer a stable job ID from process metadata."""
    pid = _safe_int(process.get("pid")) or 0
    ppid = _safe_int(process.get("ppid")) or 0
    name = str(process.get("name", "")).strip().lower()
    cmdline = str(process.get("cmdline", process.get("cmdline_summary", "")))

    run_id = _first_match(_RUN_ID_PATTERNS, cmdline)
    if run_id:
        return f"run-{_slug(run_id)}"

    master_addr = _first_match((_MASTER_PATTERN.pattern,), cmdline)
    if master_addr:
        script_match = _SCRIPT_PATTERN.search(cmdline)
        script_name = (
            script_match.group(1).split("/")[-1]
            if script_match
            else name or f"pid-{pid}"
        )
        return f"dist-{_slug(master_addr)}-{_slug(script_name)}"

    script_match = _SCRIPT_PATTERN.search(cmdline)
    if script_match and ppid > 1:
        return f"ppid-{ppid}-{_slug(script_match.group(1).split('/')[-1])}"

    if ppid > 1 and name:
        return f"ppid-{ppid}-{_slug(name)}"

    return f"pid-{pid}"


def _infer_role(process: Dict[str, Any], rank: Optional[int]) -> str:
    """Infer process role within a job."""
    cmdline = str(process.get("cmdline", process.get("cmdline_summary", ""))).lower()
    name = str(process.get("name", "")).lower()

    if any(hint in cmdline for hint in _LAUNCHER_HINTS) or any(
        hint in name for hint in ("torchrun", "deepspeed", "accelerate")
    ):
        return "launcher"
    if rank is not None:
        return "worker"
    return "process"


def _extract_rank(cmdline: str) -> Optional[int]:
    """Extract distributed rank from command line."""
    return _safe_int(_first_match(_RANK_PATTERNS, cmdline))


def _extract_world_size(cmdline: str) -> Optional[int]:
    """Extract distributed world-size from command line."""
    return _safe_int(_first_match(_WORLD_SIZE_PATTERNS, cmdline))


def _gpu_util_from_processes(gpu_processes: List[Dict[str, Any]]) -> Dict[int, float]:
    """Build PID -> GPU utilization map from vendor process metadata."""
    by_pid: Dict[int, float] = {}
    for gpu_proc in gpu_processes:
        pid = _safe_int(gpu_proc.get("pid"))
        if pid is None:
            continue

        candidates = [
            gpu_proc.get("gpu_util"),
            gpu_proc.get("sm_util"),
            gpu_proc.get("cu_occupancy"),
            gpu_proc.get("mem_util"),
        ]
        best = max(_safe_float(candidate) for candidate in candidates)
        by_pid[pid] = max(by_pid.get(pid, 0.0), best)
    return by_pid


def annotate_job_topology(
    processes: List[Dict[str, Any]],
    gpu_processes: Optional[List[Dict[str, Any]]] = None,
) -> List[Dict[str, Any]]:
    """Annotate processes with inferred job topology and straggler metadata."""
    if not processes:
        return processes

    gpu_util_by_pid = _gpu_util_from_processes(gpu_processes or [])
    grouped: Dict[str, List[Dict[str, Any]]] = {}

    for process in processes:
        cmdline = str(process.get("cmdline", process.get("cmdline_summary", "")))
        pid = _safe_int(process.get("pid")) or 0
        rank = _extract_rank(cmdline)
        world_size = _extract_world_size(cmdline)
        job_id = _infer_job_id(process)
        role = _infer_role(process, rank)

        cpu_percent = _safe_float(process.get("cpu_percent", 0.0))
        gpu_percent = max(
            _safe_float(process.get("gpu_util", 0.0)),
            gpu_util_by_pid.get(pid, 0.0),
        )
        activity_score = (0.6 * gpu_percent + 0.4 * cpu_percent) / 100.0

        process["job_id"] = job_id
        process["job_rank"] = rank
        process["job_world_size"] = world_size
        process["job_role"] = role
        process["job_activity_score"] = round(activity_score, 4)
        process["gpu_util"] = gpu_percent
        grouped.setdefault(job_id, []).append(process)

    for job_processes in grouped.values():
        scores = [float(proc.get("job_activity_score", 0.0)) for proc in job_processes]
        job_size = len(job_processes)
        peak_score = max(scores) if scores else 0.0
        median_score = median(scores) if scores else 0.0
        distributed = job_size > 1 or any(
            (proc.get("job_world_size") or 0) > 1 for proc in job_processes
        )

        straggler_count = 0
        for process in job_processes:
            score = float(process.get("job_activity_score", 0.0))
            # Conservative by design: avoid noisy straggler labels on idle jobs.
            is_straggler = (
                distributed
                and peak_score >= 0.25
                and score <= max(0.12, median_score * 0.5)
                and (peak_score - score) >= 0.25
            )

            if is_straggler:
                straggler_count += 1

            process["job_size"] = job_size
            process["is_distributed_job"] = distributed
            process["is_job_straggler"] = is_straggler
            process["job_straggler_score"] = round(
                max(0.0, peak_score - score) * 100, 1
            )

        for process in job_processes:
            process["job_straggler_count"] = straggler_count

    return processes


def summarize_job_topology(
    processes: List[Dict[str, Any]],
    max_jobs: int = 5,
) -> List[Dict[str, Any]]:
    """Build concise job summaries for UI display."""
    if not processes:
        return []

    grouped: Dict[str, List[Dict[str, Any]]] = {}
    for process in processes:
        job_id = str(process.get("job_id", "")).strip()
        if not job_id:
            continue
        grouped.setdefault(job_id, []).append(process)

    summaries: List[Dict[str, Any]] = []
    for job_id, job_processes in grouped.items():
        size = len(job_processes)
        stragglers = sum(1 for proc in job_processes if proc.get("is_job_straggler"))
        avg_cpu = (
            sum(_safe_float(proc.get("cpu_percent")) for proc in job_processes) / size
        )
        avg_gpu = (
            sum(_safe_float(proc.get("gpu_util")) for proc in job_processes) / size
        )

        summaries.append(
            {
                "job_id": job_id,
                "size": size,
                "stragglers": stragglers,
                "is_distributed": any(
                    bool(proc.get("is_distributed_job")) for proc in job_processes
                ),
                "avg_cpu": round(avg_cpu, 1),
                "avg_gpu": round(avg_gpu, 1),
            }
        )

    summaries.sort(
        key=lambda item: (
            item["stragglers"],
            item["avg_gpu"],
            item["avg_cpu"],
            item["size"],
        ),
        reverse=True,
    )
    return summaries[: max(0, max_jobs)]


__all__ = ["annotate_job_topology", "summarize_job_topology"]
