import heapq
import itertools
import math
from collections import defaultdict, deque

import numpy as np

from .common import sample_lognormal


def simulate_task_batch(task_type, task_count, params, rng, timeout_probability=None,
                        max_retries=2, include_context=True):
    task_params = params["task_types"][task_type] if include_context else _globalized(params)
    rows = []
    for _ in range(task_count):
        split = rng.random() < task_params["split_probability"]
        subtasks = max(1, int(rng.poisson(max(task_params["mean_subtasks"] - 1.0, 0.01)) + 1)) if split else 1
        subtasks = min(subtasks, 8)
        input_size = sample_lognormal(rng, task_params["external_size_bytes"]["mean"],
                                      task_params["external_size_bytes"]["cv"])
        total_bytes = 0.0
        messages = 0
        tool_calls = 0
        retries_total = 0
        success = True
        durations = []
        coordinator_think = sample_lognormal(rng, task_params["think_time_ms"]["mean"],
                                              task_params["think_time_ms"]["cv"])
        durations.append(coordinator_think)
        for _sub in range(subtasks):
            tool = rng.random() < task_params["tool_probability"]
            req = sample_lognormal(rng, task_params["request_size_bytes"]["mean"],
                                   task_params["request_size_bytes"]["cv"])
            total_bytes += req
            messages += 1
            worker_time = sample_lognormal(rng, task_params["think_time_ms"]["mean"],
                                           task_params["think_time_ms"]["cv"])
            durations.append(worker_time)
            if tool:
                p_timeout = task_params["retry_probability"] if timeout_probability is None else timeout_probability
                attempts = 0
                tool_success = False
                while attempts <= max_retries:
                    tool_calls += 1
                    tool_req = sample_lognormal(rng, task_params["request_size_bytes"]["mean"] * 0.35,
                                                task_params["request_size_bytes"]["cv"])
                    total_bytes += tool_req
                    messages += 1
                    timed_out = rng.random() < p_timeout
                    if not timed_out:
                        tool_resp = sample_lognormal(rng, task_params["response_size_bytes"]["mean"] * 0.72,
                                                     task_params["response_size_bytes"]["cv"])
                        total_bytes += tool_resp
                        messages += 1
                        durations.append(sample_lognormal(rng, task_params["think_time_ms"]["mean"] * 0.65,
                                                          task_params["think_time_ms"]["cv"]))
                        tool_success = True
                        break
                    if attempts < max_retries:
                        retries_total += 1
                        durations.append(2200.0 + 180.0 * (attempts + 1))
                    attempts += 1
                success = success and tool_success
            response = sample_lognormal(rng, task_params["response_size_bytes"]["mean"],
                                        task_params["response_size_bytes"]["cv"])
            total_bytes += response if success else 600.0
            messages += 1
        final_size = max(800.0, subtasks * task_params["response_size_bytes"]["mean"] * 0.38)
        total_bytes += final_size
        messages += 1
        latency = sum(durations) + 8.0 * messages
        rows.append({
            "task_type": task_type,
            "internal_bytes": total_bytes,
            "input_bytes": input_size,
            "amplification": total_bytes / max(input_size, 1.0),
            "latency_ms": latency,
            "message_count": messages,
            "tool_calls": tool_calls,
            "retries": retries_total,
            "success": success,
        })
    return rows


def _globalized(params):
    entries = list(params["task_types"].values())
    def mean(name):
        return float(np.mean([e[name] for e in entries]))
    def stat(name):
        return {
            "mean": float(np.mean([e[name]["mean"] for e in entries])),
            "cv": float(np.mean([e[name]["cv"] for e in entries])),
        }
    return {
        "tool_probability": mean("tool_probability"),
        "split_probability": mean("split_probability"),
        "mean_subtasks": mean("mean_subtasks"),
        "retry_probability": mean("retry_probability"),
        "failure_probability": mean("failure_probability"),
        "think_time_ms": stat("think_time_ms"),
        "external_size_bytes": stat("external_size_bytes"),
        "request_size_bytes": stat("request_size_bytes"),
        "response_size_bytes": stat("response_size_bytes"),
    }


def simulate_queue_stress(arrival_rate, duration_seconds, warmup_seconds, task_type, params,
                          rng, max_concurrency=4, queue_capacity=500, timeout_probability=None,
                          max_retries=2):
    tasks = simulate_task_batch(
        task_type, max(200, int(arrival_rate * duration_seconds * 1.5)), params, rng,
        timeout_probability=timeout_probability, max_retries=max_retries,
    )
    event_queue = []
    counter = itertools.count()
    time = 0.0
    task_index = 0
    while time < duration_seconds:
        time += float(rng.exponential(1.0 / max(arrival_rate, 1e-9)))
        if time <= duration_seconds:
            heapq.heappush(event_queue, (time, next(counter), "arrival", task_index))
            task_index += 1

    waiting = deque()
    busy = 0
    completed = []
    dropped = 0
    queue_peak = 0
    arrival_times = {}
    while event_queue:
        now, _, event_type, idx = heapq.heappop(event_queue)
        if event_type == "arrival":
            arrival_times[idx] = now
            if busy < max_concurrency:
                busy += 1
                task = tasks[idx % len(tasks)]
                service = task["latency_ms"] / 1000.0
                heapq.heappush(event_queue, (now + service, next(counter), "complete", idx))
            elif len(waiting) < queue_capacity:
                waiting.append(idx)
                queue_peak = max(queue_peak, len(waiting))
            else:
                dropped += 1
        else:
            task = tasks[idx % len(tasks)]
            if arrival_times[idx] >= warmup_seconds:
                record = dict(task)
                record["end_to_end_ms"] = (now - arrival_times[idx]) * 1000.0
                completed.append(record)
            busy -= 1
            if waiting:
                nxt = waiting.popleft()
                busy += 1
                task2 = tasks[nxt % len(tasks)]
                heapq.heappush(event_queue, (now + task2["latency_ms"] / 1000.0,
                                             next(counter), "complete", nxt))
    offered = max(1, task_index)
    if completed:
        delays = np.array([x["end_to_end_ms"] for x in completed])
        amps = np.array([x["amplification"] for x in completed])
        success = np.array([x["success"] for x in completed], dtype=float)
    else:
        delays, amps, success = np.array([0.0]), np.array([0.0]), np.array([0.0])
    return {
        "arrival_rate": arrival_rate,
        "offered_tasks": offered,
        "completed_tasks": len(completed),
        "throughput_per_second": len(completed) / max(duration_seconds - warmup_seconds, 1),
        "mean_delay_ms": float(delays.mean()),
        "p95_delay_ms": float(np.quantile(delays, 0.95)),
        "queue_peak": queue_peak,
        "drop_rate": dropped / offered,
        "success_rate": float(success.mean()),
        "mean_amplification": float(amps.mean()),
    }
