from __future__ import annotations

import math
import random
from typing import Iterable

from .entities import Agent, SystemState, Ticket


def severity_priority(severity: str) -> int:
    return {"low": 0, "medium": 1, "high": 2}.get(severity, 0)


def apply_action(ticket: Ticket, action_type: str) -> Ticket:
    updated = ticket.model_copy(deep=True)
    updated.attempts_used += 1

    if action_type in {"triage", "respond", "assign", "reprioritize"}:
        updated.status = "in_progress"
    elif action_type == "resolve":
        updated.status = "resolved"
    elif action_type == "escalate":
        updated.status = "in_progress"
        updated.priority = max(updated.priority, 2)

    if updated.attempts_used >= updated.max_attempts and updated.status != "resolved":
        updated.status = "open"

    return updated


def is_terminal(ticket: Ticket) -> bool:
    return ticket.status == "resolved" or ticket.attempts_used >= ticket.max_attempts


def check_sla(ticket: Ticket, current_time: float) -> bool:
    if ticket.deadline is None:
        return True
    return current_time <= ticket.deadline


def compute_sla_penalty(ticket: Ticket, current_time: float) -> float:
    if ticket.deadline is None:
        return 0.0
    delay = max(0.0, current_time - ticket.deadline)
    multiplier = {"low": 1.0, "medium": 2.0, "high": 3.0}[ticket.severity]
    return 100.0 * multiplier * (1.0 + delay)


def is_blocked(ticket: Ticket, completed_ids: set[str]) -> bool:
    return any(dep_id not in completed_ids for dep_id in ticket.dependencies)


def get_available_tickets(waiting_queue: Iterable[Ticket], completed_ids: set[str]) -> list[Ticket]:
    return [ticket for ticket in waiting_queue if not is_blocked(ticket, completed_ids)]


def generate_arrivals(
    current_time: float,
    time_delta: float,
    arrival_rate: float,
    rng: random.Random,
    count_seed: int,
) -> list[Ticket]:
    expected = max(arrival_rate * time_delta, 0.0)
    integer_part = int(expected)
    fractional = expected - integer_part
    arrivals = integer_part + (1 if rng.random() < fractional else 0)

    tickets: list[Ticket] = []
    for i in range(arrivals):
        severity = rng.choices(["low", "medium", "high"], weights=[0.5, 0.3, 0.2], k=1)[0]
        ticket_id = f"A-{int(current_time)}-{count_seed}-{i}"
        tickets.append(
            Ticket(
                id=ticket_id,
                summary=f"Auto-generated ticket at t={current_time:.1f}",
                severity=severity,
                created_at=current_time + rng.random() * max(time_delta, 1e-6),
                deadline=current_time + 6.0,
                priority=severity_priority(severity),
            )
        )
    return tickets


def sample_agent_failure(agent: Agent, time_delta: float, rng: random.Random) -> tuple[bool, float]:
    base_prob = 0.01
    load_ratio = agent.current_load / max(agent.capacity, 1e-6)
    stress = load_ratio**2 if load_ratio > 1.0 else 1.0
    prob = base_prob * stress * max(time_delta, 0.0)

    if rng.random() < prob:
        return True, rng.uniform(1.0, 4.0)
    return False, 0.0


def advance_time(state: SystemState, delta: float) -> None:
    state.current_time += delta
    active_agent_cost = sum(agent.hourly_cost for agent in state.agents if state.current_time >= agent.failed_until)
    state.total_cost += active_agent_cost * delta


def pick_next_ticket(state: SystemState) -> Ticket | None:
    if not state.waiting_queue:
        return None

    completed_ids = {ticket.id for ticket in state.completed_tickets}
    available = get_available_tickets(state.waiting_queue, completed_ids)
    if not available:
        return None

    available.sort(key=lambda ticket: (ticket.priority, severity_priority(ticket.severity), -ticket.created_at), reverse=True)
    selected = available[0]
    state.waiting_queue = [ticket for ticket in state.waiting_queue if ticket.id != selected.id]
    return selected