"""Lockstep batched branch decoding for the Transformers backend."""

from __future__ import annotations

from collections import deque
from collections.abc import Iterator
from dataclasses import dataclass, field
from typing import Any, Protocol

import torch
import torch.nn.functional as F


class _Engine(Protocol):
    model: Any
    processor: Any
    tokenizer: Any
    device: torch.device
    eos_token_id: int
    fork_token_map: dict[int, list[int]]
    max_new_tokens: int
    max_branch_tokens: int
    max_concurrent_branches: int
    max_total_branches: int
    strict: bool
    _branch_forbidden_ids: list[int]

    def _position_ids(
        self,
        logical_positions: torch.LongTensor,
        *,
        multimodal: bool,
    ) -> torch.LongTensor: ...

    def _next_token(self, logits: torch.Tensor, *, branch: bool) -> int: ...


CacheState = list[dict[str, Any]]


@dataclass(slots=True)
class _Stream:
    past: CacheState
    cache_len: int
    next_input: int | None
    branch_index: int | None = None
    fork_position: int | None = None
    injected_tokens: list[int] = field(default_factory=list)
    generated_tokens: list[int] = field(default_factory=list)
    state: str = "active"


def _cache_to_state(cache: Any) -> CacheState:
    """Extract Qwen DynamicCache tensors into independently batchable state."""
    try:
        from transformers.cache_utils import LinearAttentionLayer
    except ImportError as exc:  # pragma: no cover - pinned Transformers provides it
        raise RuntimeError("Parallel decoding requires Transformers DynamicCache support.") from exc

    if not hasattr(cache, "layers"):
        raise TypeError(
            "Parallel decoding requires a Transformers DynamicCache; "
            f"received {type(cache).__name__}."
        )

    state: CacheState = []
    for layer in cache.layers:
        if isinstance(layer, LinearAttentionLayer):
            if not layer.is_conv_states_initialized or not layer.is_recurrent_states_initialized:
                raise RuntimeError("A linear-attention cache layer was not initialized by prefill.")
            state.append(
                {
                    "kind": "linear",
                    "conv": layer.conv_states.contiguous(),
                    "recur": layer.recurrent_states.contiguous(),
                    "has_previous_state": bool(layer.has_previous_state),
                }
            )
            continue

        keys = getattr(layer, "keys", None)
        values = getattr(layer, "values", None)
        if keys is None or values is None:
            raise TypeError(
                "Parallel decoding only supports initialized dynamic or linear cache layers; "
                f"received {type(layer).__name__}."
            )
        state.append(
            {
                "kind": "full",
                "keys": keys.contiguous(),
                "values": values.contiguous(),
            }
        )
    return state


def _clone_state(state: CacheState, *, truncate_to: int | None = None) -> CacheState:
    cloned: CacheState = []
    for layer in state:
        if layer["kind"] == "full":
            keys = layer["keys"]
            values = layer["values"]
            if truncate_to is not None:
                keys = keys[:, :, :truncate_to, :]
                values = values[:, :, :truncate_to, :]
            cloned.append(
                {
                    "kind": "full",
                    "keys": keys.contiguous().clone(),
                    "values": values.contiguous().clone(),
                }
            )
        else:
            cloned.append(
                {
                    "kind": "linear",
                    "conv": layer["conv"].contiguous().clone(),
                    "recur": layer["recur"].contiguous().clone(),
                    "has_previous_state": layer["has_previous_state"],
                }
            )
    return cloned


def _state_to_cache(state: CacheState):
    from transformers.cache_utils import DynamicCache, DynamicLayer, LinearAttentionLayer

    cache = DynamicCache()
    layers = []
    for entry in state:
        if entry["kind"] == "full":
            layer = DynamicLayer()
            layer.update(entry["keys"], entry["values"])
        else:
            layer = LinearAttentionLayer()
            conv = entry["conv"]
            recur = entry["recur"]
            layer.lazy_initialization(conv_states=conv, recurrent_states=recur)
            layer.conv_states.copy_(conv)
            layer.recurrent_states.copy_(recur)
            layer.has_previous_state = entry["has_previous_state"]
        layers.append(layer)
    cache.layers = layers
    return cache


def _batch_cache(streams: list[_Stream], max_len: int):
    from transformers.cache_utils import DynamicCache, DynamicLayer, LinearAttentionLayer

    cache = DynamicCache()
    layers = []
    for layer_index in range(len(streams[0].past)):
        first = streams[0].past[layer_index]
        if first["kind"] == "full":
            keys = []
            values = []
            for stream in streams:
                entry = stream.past[layer_index]
                pad = max_len - stream.cache_len
                keys.append(F.pad(entry["keys"], (0, 0, pad, 0)))
                values.append(F.pad(entry["values"], (0, 0, pad, 0)))
            layer = DynamicLayer()
            layer.update(torch.cat(keys, dim=0), torch.cat(values, dim=0))
        else:
            conv = torch.cat([stream.past[layer_index]["conv"] for stream in streams], dim=0)
            recur = torch.cat(
                [stream.past[layer_index]["recur"] for stream in streams], dim=0
            )
            layer = LinearAttentionLayer()
            layer.lazy_initialization(conv_states=conv, recurrent_states=recur)
            layer.conv_states.copy_(conv)
            layer.recurrent_states.copy_(recur)
            layer.has_previous_state = any(
                stream.past[layer_index]["has_previous_state"] for stream in streams
            )
        layers.append(layer)
    cache.layers = layers
    return cache


def _split_batch_cache(cache: Any, streams: list[_Stream]) -> None:
    from transformers.cache_utils import LinearAttentionLayer

    new_lengths = [stream.cache_len + 1 for stream in streams]
    for batch_index, stream in enumerate(streams):
        state: CacheState = []
        for layer in cache.layers:
            if isinstance(layer, LinearAttentionLayer):
                state.append(
                    {
                        "kind": "linear",
                        "conv": layer.conv_states[batch_index : batch_index + 1].contiguous(),
                        "recur": layer.recurrent_states[
                            batch_index : batch_index + 1
                        ].contiguous(),
                        "has_previous_state": bool(layer.has_previous_state),
                    }
                )
                continue
            length = new_lengths[batch_index]
            state.append(
                {
                    "kind": "full",
                    "keys": layer.keys[
                        batch_index : batch_index + 1, :, -length:, :
                    ].contiguous(),
                    "values": layer.values[
                        batch_index : batch_index + 1, :, -length:, :
                    ].contiguous(),
                }
            )
        stream.past = state
        stream.cache_len = new_lengths[batch_index]


def _scheduler_event(
    main: _Stream,
    branches: list[_Stream],
    *,
    phase: str,
    batch_size: int = 0,
) -> dict[str, Any]:
    return {
        "type": "scheduler",
        "execution_mode": "parallel",
        "phase": phase,
        "main_active": main.state == "active",
        "active_branches": sum(branch.state == "active" for branch in branches),
        "queued_branches": sum(branch.state == "queued" for branch in branches),
        "completed_branches": sum(branch.state == "done" for branch in branches),
        "batch_size": batch_size,
    }


def stream_parallel(engine: _Engine, messages: list[dict[str, Any]]) -> Iterator[dict[str, Any]]:
    """Decode main and active branches together in a per-step GPU batch."""
    with torch.inference_mode():
        yield from _stream_parallel(engine, messages)


def _stream_parallel(engine: _Engine, messages: list[dict[str, Any]]) -> Iterator[dict[str, Any]]:
    yield {"type": "accepted", "execution_mode": "parallel"}

    inputs = engine.processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(engine.device)
    prompt_ids = inputs["input_ids"][0].tolist()
    prompt_length = len(prompt_ids)
    multimodal = inputs.get("image_grid_thw") is not None or inputs.get("video_grid_thw") is not None

    output = engine.model(**inputs, use_cache=True)
    first_token = engine._next_token(output.logits, branch=False)
    main = _Stream(
        past=_cache_to_state(output.past_key_values),
        cache_len=prompt_length,
        next_input=first_token,
        generated_tokens=[first_token],
    )
    del output

    if first_token == engine.eos_token_id or engine.max_new_tokens == 1:
        main.state = "done"

    branches: list[_Stream] = []
    pending: deque[int] = deque()
    active: set[int] = set()
    dropped_triggers = 0
    peak_batch_size = 1

    yield _scheduler_event(main, branches, phase="prefill_complete", batch_size=1)
    yield {
        "type": "main",
        "token_ids": [first_token],
        "delta_text": engine.tokenizer.decode([first_token], skip_special_tokens=False),
        "total": 1,
    }
    if main.state == "done":
        yield {"type": "main_done", "total": 1}

    def register_fork(token_id: int) -> dict[str, Any] | None:
        nonlocal dropped_triggers
        injected = engine.fork_token_map.get(token_id)
        if not injected:
            return None
        if len(branches) >= engine.max_total_branches:
            dropped_triggers += 1
            return None

        fork_position = main.cache_len
        expected_position = prompt_length + len(main.generated_tokens) - 1
        if fork_position != expected_position:
            raise RuntimeError(
                "Main cache is not aligned with its fork token: "
                f"cache={fork_position}, token={expected_position}."
            )
        branch_index = len(branches)
        injected_tokens = list(injected)
        if not injected_tokens:
            raise ValueError("A fork target must contain at least one token.")
        branches.append(
            _Stream(
                past=_clone_state(main.past, truncate_to=fork_position),
                cache_len=fork_position,
                next_input=None,
                branch_index=branch_index,
                fork_position=fork_position,
                injected_tokens=injected_tokens,
                state="queued",
            )
        )
        pending.append(branch_index)
        return {
            "type": "fork",
            "branch_index": branch_index,
            "fork_position": fork_position,
            "trigger_token_id": token_id,
            "injected_token_ids": injected_tokens,
            "injected_text": engine.tokenizer.decode(
                injected_tokens, skip_special_tokens=False
            ),
            "branch_state": "queued",
        }

    fork_event = register_fork(first_token)
    if fork_event is not None:
        yield fork_event

    def activate_pending() -> Iterator[dict[str, Any]]:
        while pending and len(active) < engine.max_concurrent_branches:
            branch_index = pending.popleft()
            branch = branches[branch_index]
            branch.state = "active"
            active.add(branch_index)

            injected = torch.tensor(
                [branch.injected_tokens], dtype=torch.long, device=engine.device
            )
            logical = torch.arange(
                branch.fork_position,
                branch.fork_position + len(branch.injected_tokens),
                dtype=torch.long,
                device=engine.device,
            ).unsqueeze(0)
            output = engine.model(
                input_ids=injected,
                position_ids=engine._position_ids(logical, multimodal=multimodal),
                past_key_values=_state_to_cache(branch.past),
                use_cache=True,
            )
            branch.past = _cache_to_state(output.past_key_values)
            branch.cache_len += len(branch.injected_tokens)
            first_branch_token = engine._next_token(output.logits, branch=True)
            branch.generated_tokens.append(first_branch_token)
            branch.next_input = first_branch_token
            del output

            yield {
                "type": "branch",
                "branch_index": branch_index,
                "fork_position": branch.fork_position,
                "token_ids": [first_branch_token],
                "delta_text": engine.tokenizer.decode(
                    [first_branch_token], skip_special_tokens=False
                ),
                "total": len(branch.injected_tokens) + 1,
            }
            if (
                first_branch_token == engine.eos_token_id
                or len(branch.generated_tokens) >= engine.max_branch_tokens
            ):
                branch.state = "done"
                active.discard(branch_index)
                yield {
                    "type": "branch_done",
                    "branch_index": branch_index,
                    "total": len(branch.injected_tokens) + len(branch.generated_tokens),
                }
            yield _scheduler_event(main, branches, phase="branch_started")

    yield from activate_pending()

    while main.state == "active" or active or pending:
        yield from activate_pending()
        streams: list[_Stream] = []
        if main.state == "active":
            streams.append(main)
        streams.extend(branches[index] for index in sorted(active))
        if not streams:
            continue

        batch_size = len(streams)
        peak_batch_size = max(peak_batch_size, batch_size)
        yield _scheduler_event(main, branches, phase="decoding", batch_size=batch_size)

        max_len = max(stream.cache_len for stream in streams)
        input_ids = torch.tensor(
            [[stream.next_input] for stream in streams],
            dtype=torch.long,
            device=engine.device,
        )
        logical_positions = torch.tensor(
            [[stream.cache_len] for stream in streams],
            dtype=torch.long,
            device=engine.device,
        )
        attention_mask = torch.zeros(
            len(streams), max_len + 1, dtype=torch.long, device=engine.device
        )
        for row, stream in enumerate(streams):
            attention_mask[row, max_len - stream.cache_len :] = 1

        output = engine.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=engine._position_ids(logical_positions, multimodal=multimodal),
            past_key_values=_batch_cache(streams, max_len),
            use_cache=True,
        )
        logits = output.logits[:, -1, :]
        if engine.strict and engine._branch_forbidden_ids:
            logits = logits.clone()
            forbidden = torch.tensor(
                engine._branch_forbidden_ids, dtype=torch.long, device=logits.device
            )
            for row, stream in enumerate(streams):
                if stream.branch_index is not None:
                    logits[row, forbidden] = torch.finfo(logits.dtype).min
        token_ids = logits.argmax(dim=-1).cpu().tolist()
        _split_batch_cache(output.past_key_values, streams)
        del output, logits

        new_forks: list[dict[str, Any]] = []
        for stream, token_id_raw in zip(streams, token_ids, strict=True):
            token_id = int(token_id_raw)
            stream.generated_tokens.append(token_id)
            stream.next_input = token_id
            if stream.branch_index is None:
                yield {
                    "type": "main",
                    "token_ids": [token_id],
                    "delta_text": engine.tokenizer.decode(
                        [token_id], skip_special_tokens=False
                    ),
                    "total": len(main.generated_tokens),
                }
                fork_event = register_fork(token_id)
                if fork_event is not None:
                    new_forks.append(fork_event)
                if (
                    token_id == engine.eos_token_id
                    or len(main.generated_tokens) >= engine.max_new_tokens
                ):
                    main.state = "done"
                    yield {"type": "main_done", "total": len(main.generated_tokens)}
                continue

            branch_index = stream.branch_index
            yield {
                "type": "branch",
                "branch_index": branch_index,
                "fork_position": stream.fork_position,
                "token_ids": [token_id],
                "delta_text": engine.tokenizer.decode([token_id], skip_special_tokens=False),
                "total": len(stream.injected_tokens) + len(stream.generated_tokens),
            }
            if (
                token_id == engine.eos_token_id
                or len(stream.generated_tokens) >= engine.max_branch_tokens
            ):
                stream.state = "done"
                active.discard(branch_index)
                yield {
                    "type": "branch_done",
                    "branch_index": branch_index,
                    "total": len(stream.injected_tokens) + len(stream.generated_tokens),
                }

        yield from new_forks

    result_branches = []
    for branch in branches:
        token_ids = branch.injected_tokens + branch.generated_tokens
        result_branches.append(
            {
                "branch_index": branch.branch_index,
                "fork_position": branch.fork_position,
                "injected_token_ids": branch.injected_tokens,
                "token_ids": token_ids,
                "text": engine.tokenizer.decode(token_ids, skip_special_tokens=False),
            }
        )

    yield {
        "type": "done",
        "execution_mode": "parallel",
        "main": engine.tokenizer.decode(main.generated_tokens, skip_special_tokens=False),
        "main_token_ids": main.generated_tokens,
        "prompt_tokens": prompt_length,
        "branches": result_branches,
        "branch_limit_reached": dropped_triggers > 0,
        "dropped_branch_triggers": dropped_triggers,
        "peak_batch_size": peak_batch_size,
    }