import enum
from dataclasses import dataclass, field
from typing import List, Dict, Optional
import time
import heapq


class ProcessState(enum.Enum):
    """操作系统中的进程状态"""
    NEW = 0  # 进程刚创建
    READY = 1  # 等待被分配到处理器的进程
    RUNNING = 2  # 当前正在执行的进程
    WAITING = 3  # 等待某些事件的进程（I/O、资源分配）
    TERMINATED = 4  # 已完成执行的进程


@dataclass
class Process:
    """表示操作系统中的进程"""
    pid: int
    script_path: str
    state: ProcessState = field(default=ProcessState.NEW)
    priority: int = field(default=0)
    arrival_time: float = field(default_factory=time.time)
    start_time: Optional[float] = None
    completion_time: Optional[float] = None
    cpu_time: float = 0.0
    waiting_time: float = 0.0


class ProcessManager:
    """
    管理进程调度和生命周期
    实现多种调度算法
    """

    def __init__(self):
        self.processes: Dict[int, Process] = {}
        self.ready_queue: List[Process] = []
        self.next_pid = 1
        self.total_processes = 0

    def create_process(self, script_path: str, priority: int = 0) -> Process:
        """
        创建新进程

        :param script_path: 进程脚本路径
        :param priority: 进程优先级（数字越小优先级越高）
        :return: 创建的进程
        """
        process = Process(
            pid=self.next_pid,
            script_path=script_path,
            priority=priority
        )
        self.processes[self.next_pid] = process
        self.next_pid += 1
        self.total_processes += 1

        # 将进程移至就绪状态
        self._change_process_state(process, ProcessState.READY)

        return process

    def _change_process_state(self, process: Process, new_state: ProcessState):
        """
        改变进程状态并更新相关时间戳

        :param process: 要更新的进程
        :param new_state: 进程的新状态
        """
        old_state = process.state
        process.state = new_state

        # 跟踪性能指标的计时
        current_time = time.time()
        if new_state == ProcessState.RUNNING:
            process.start_time = current_time
        elif new_state == ProcessState.TERMINATED:
            process.completion_time = current_time

    def schedule_process(self) -> Optional[Process]:
        """
        选择并调度要运行的进程
        使用基于优先级的抢占式调度

        :return: 要运行的进程，如果没有就绪进程则返回None
        """
        if not self.ready_queue:
            return None

        # 按优先级对就绪队列排序（数字越小优先级越高）
        self.ready_queue.sort(key=lambda p: p.priority)

        # 选择最高优先级进程
        selected_process = self.ready_queue[0]
        self._change_process_state(selected_process, ProcessState.RUNNING)

        return selected_process

    def terminate_process(self, process: Process):
        """
        终止进程并更新其状态

        :param process: 要终止的进程
        """
        self._change_process_state(process, ProcessState.TERMINATED)

        # 如果在就绪队列中，则移除
        if process in self.ready_queue:
            self.ready_queue.remove(process)

    def get_total_processes(self) -> int:
        """
        获取已创建的进程总数

        :return: 进程总数
        """
        return self.total_processes

    def calculate_avg_turnaround_time(self) -> float:
        """
        计算已完成进程的平均周转时间

        :return: 平均周转时间
        """
        completed_processes = [
            p for p in self.processes.values()
            if p.completion_time is not None
        ]

        if not completed_processes:
            return 0.0

        turnaround_times = [
            p.completion_time - p.arrival_time
            for p in completed_processes
        ]

        return sum(turnaround_times) / len(turnaround_times)

    def calculate_avg_weighted_turnaround_time(self) -> float:
        """
        计算平均带权周转时间
        带权周转时间 = (完成时间 - 到达时间) / CPU时间

        :return: 平均带权周转时间
        """
        completed_processes = [
            p for p in self.processes.values()
            if p.completion_time is not None and p.cpu_time > 0
        ]

        if not completed_processes:
            return 0.0

        weighted_turnaround_times = [
            (p.completion_time - p.arrival_time) / p.cpu_time
            for p in completed_processes
        ]

        return sum(weighted_turnaround_times) / len(weighted_turnaround_times)