"""Transformers inference for sequential and batched PaDoc decoding."""

from __future__ import annotations

import argparse
import copy
import json
import logging
from collections.abc import Iterator
from pathlib import Path
from typing import Any

import torch

from padoc.modeling import load_padoc_model

logger = logging.getLogger(__name__)

EXECUTION_MODES = ("parallel", "sequential")


def _fork_id_map(tokenizer, values: dict[str, str]) -> dict[int, list[int]]:
    result: dict[int, list[int]] = {}
    for trigger, target in values.items():
        trigger_ids = tokenizer.encode(trigger, add_special_tokens=False)
        target_ids = tokenizer.encode(target, add_special_tokens=False)
        if len(trigger_ids) != 1 or len(target_ids) != 1:
            raise ValueError(
                f"Invalid atomic fork mapping {trigger!r}->{target!r}: {trigger_ids}->{target_ids}"
            )
        result[trigger_ids[0]] = list(target_ids)
    return result


class SequentialPaDocEngine:
    """Greedy PaDoc decoder with sequential and lockstep-batched execution."""

    def __init__(
        self,
        model,
        processor,
        fork_token_map: dict[str, str] | dict[int, list[int]],
        *,
        max_new_tokens: int = 512,
        max_branch_tokens: int | None = None,
        max_concurrent_branches: int = 8,
        max_total_branches: int = 64,
        execution_mode: str = "sequential",
        strict: bool = True,
    ) -> None:
        if max_new_tokens < 1 or (max_branch_tokens is not None and max_branch_tokens < 1):
            raise ValueError("Token limits must be positive.")
        if max_total_branches < 0:
            raise ValueError("max_total_branches must be non-negative.")
        if max_concurrent_branches < 1:
            raise ValueError("max_concurrent_branches must be positive.")
        if execution_mode not in EXECUTION_MODES:
            raise ValueError(
                f"Unknown execution_mode {execution_mode!r}; choose from {EXECUTION_MODES}."
            )
        self.model = model
        self.processor = processor
        self.tokenizer = processor.tokenizer
        if fork_token_map and isinstance(next(iter(fork_token_map)), str):
            fork_token_map = _fork_id_map(self.tokenizer, fork_token_map)
        self.fork_token_map = dict(fork_token_map)
        self.max_new_tokens = max_new_tokens
        self.max_branch_tokens = max_branch_tokens or max_new_tokens
        self.max_concurrent_branches = (
            min(max_concurrent_branches, max_total_branches)
            if max_total_branches
            else max_concurrent_branches
        )
        self.max_total_branches = max_total_branches
        self.execution_mode = execution_mode
        self.strict = strict
        self.eos_token_id = self.tokenizer.eos_token_id
        if self.eos_token_id is None:
            raise ValueError("The tokenizer must define eos_token_id.")
        self.device = model.device
        parameters = getattr(model, "parameters", None)
        self.devices = (
            sorted({str(parameter.device) for parameter in parameters()})
            if callable(parameters)
            else [str(self.device)]
        )

        self._branch_forbidden_ids: list[int] = []
        for token in ("<SP_LAYOUT>", "</SP_LAYOUT>"):
            token_ids = self.tokenizer.encode(token, add_special_tokens=False)
            if len(token_ids) == 1:
                self._branch_forbidden_ids.append(token_ids[0])

    @classmethod
    def from_pretrained(
        cls,
        model_path: str | Path,
        *,
        dtype: torch.dtype = torch.bfloat16,
        device: str = "cuda:0",
        attn_implementation: str = "sdpa",
        **engine_kwargs,
    ) -> SequentialPaDocEngine:
        if device == "auto":
            raise ValueError(
                "PaDoc Transformers inference is single-device; choose an explicit "
                "device such as 'cuda:0' or 'cpu'."
            )
        device_map: dict[str, str] = {"": device}
        model, processor, fork_map = load_padoc_model(
            model_path,
            dtype=dtype,
            device_map=device_map,
            attn_implementation=attn_implementation,
        )
        model.eval()
        return cls(model, processor, fork_map, **engine_kwargs)

    def resolve_execution_mode(self, execution_mode: str | None) -> str:
        mode = execution_mode or self.execution_mode
        if mode not in EXECUTION_MODES:
            raise ValueError(f"Unknown execution_mode {mode!r}; choose from {EXECUTION_MODES}.")
        return mode

    def _mrope_delta(self, *, multimodal: bool) -> int:
        if not multimodal:
            return 0
        rope_deltas = getattr(getattr(self.model, "model", None), "rope_deltas", None)
        if rope_deltas is None or rope_deltas.numel() == 0:
            raise RuntimeError("Multimodal prefill did not produce M-RoPE deltas.")
        return int(rope_deltas.reshape(-1)[0].item())

    def _position_ids(
        self,
        logical_positions: torch.LongTensor,
        *,
        multimodal: bool,
    ) -> torch.LongTensor:
        shifted = logical_positions + self._mrope_delta(multimodal=multimodal)
        axes = 3 if multimodal else 4
        return shifted.unsqueeze(0).expand(axes, -1, -1).contiguous()

    def _next_token(self, logits: torch.Tensor, *, branch: bool) -> int:
        values = logits[:, -1, :]
        if branch and self.strict and self._branch_forbidden_ids:
            values = values.clone()
            values[:, self._branch_forbidden_ids] = torch.finfo(values.dtype).min
        return int(values.argmax(dim=-1).item())

    def _stream_branch_tokens(
        self,
        parent_cache,
        *,
        fork_position: int,
        injected_tokens: list[int],
        multimodal: bool,
    ) -> Iterator[int]:
        if not injected_tokens:
            raise ValueError("A fork target must contain at least one token.")
        cache = copy.deepcopy(parent_cache)
        injected = torch.tensor([injected_tokens], dtype=torch.long, device=self.device)
        logical = torch.arange(
            fork_position,
            fork_position + len(injected_tokens),
            dtype=torch.long,
            device=self.device,
        ).unsqueeze(0)
        output = self.model(
            input_ids=injected,
            position_ids=self._position_ids(logical, multimodal=multimodal),
            past_key_values=cache,
            use_cache=True,
        )
        cache = output.past_key_values
        next_token = self._next_token(output.logits, branch=True)
        generated = 0

        while generated < self.max_branch_tokens:
            yield next_token
            generated += 1
            if next_token == self.eos_token_id:
                break
            logical_position = fork_position + len(injected_tokens) + generated - 1
            output = self.model(
                input_ids=torch.tensor([[next_token]], dtype=torch.long, device=self.device),
                position_ids=self._position_ids(
                    torch.tensor([[logical_position]], dtype=torch.long, device=self.device),
                    multimodal=multimodal,
                ),
                past_key_values=cache,
                use_cache=True,
            )
            cache = output.past_key_values
            next_token = self._next_token(output.logits, branch=True)

    def stream(
        self,
        messages: list[dict[str, Any]],
        *,
        execution_mode: str | None = None,
    ) -> Iterator[dict[str, Any]]:
        """Yield a shared event protocol from either Transformers execution mode."""
        mode = self.resolve_execution_mode(execution_mode)
        if mode == "parallel":
            from padoc.transformers_parallel import stream_parallel

            yield from stream_parallel(self, messages)
            return
        yield from self._stream_sequential(messages)

    @torch.inference_mode()
    def _stream_sequential(self, messages: list[dict[str, Any]]) -> Iterator[dict[str, Any]]:
        """Decode a branch to completion before resuming the batch=1 main stream."""
        yield {"type": "accepted", "execution_mode": "sequential"}
        inputs = self.processor.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt",
        ).to(self.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 = self.model(**inputs, use_cache=True)
        cache = output.past_key_values
        next_token = self._next_token(output.logits, branch=False)
        main_tokens: list[int] = []
        branches: list[dict[str, Any]] = []
        dropped_triggers = 0
        yield {
            "type": "scheduler",
            "execution_mode": "sequential",
            "phase": "prefill_complete",
            "main_active": True,
            "active_branches": 0,
            "queued_branches": 0,
            "completed_branches": 0,
            "batch_size": 1,
        }

        while len(main_tokens) < self.max_new_tokens:
            response_position = len(main_tokens)
            main_tokens.append(next_token)
            yield {
                "type": "main",
                "token_ids": [next_token],
                "delta_text": self.tokenizer.decode([next_token], skip_special_tokens=False),
                "text": self.tokenizer.decode(main_tokens, skip_special_tokens=False),
                "total": len(main_tokens),
            }
            injected = self.fork_token_map.get(next_token)
            if injected:
                if len(branches) >= self.max_total_branches:
                    dropped_triggers += 1
                else:
                    fork_position = prompt_length + response_position
                    branch_index = len(branches)
                    injected_tokens = list(injected)
                    injected_text = self.tokenizer.decode(
                        injected_tokens, skip_special_tokens=False
                    )
                    yield {
                        "type": "fork",
                        "branch_index": branch_index,
                        "fork_position": fork_position,
                        "trigger_token_id": next_token,
                        "injected_token_ids": injected_tokens,
                        "injected_text": injected_text,
                        "branch_state": "active",
                    }
                    yield {
                        "type": "scheduler",
                        "execution_mode": "sequential",
                        "phase": "branch_decoding",
                        "main_active": False,
                        "active_branches": 1,
                        "queued_branches": 0,
                        "completed_branches": len(branches),
                        "batch_size": 1,
                    }
                    generated_branch_tokens: list[int] = []
                    for branch_token in self._stream_branch_tokens(
                        cache,
                        fork_position=fork_position,
                        injected_tokens=injected_tokens,
                        multimodal=multimodal,
                    ):
                        generated_branch_tokens.append(branch_token)
                        branch_tokens = [*injected_tokens, *generated_branch_tokens]
                        yield {
                            "type": "branch",
                            "branch_index": branch_index,
                            "fork_position": fork_position,
                            "token_ids": [branch_token],
                            "delta_text": self.tokenizer.decode(
                                [branch_token], skip_special_tokens=False
                            ),
                            "text": self.tokenizer.decode(
                                branch_tokens, skip_special_tokens=False
                            ),
                            "total": len(branch_tokens),
                        }
                    branch_tokens = [*injected_tokens, *generated_branch_tokens]
                    branches.append(
                        {
                            "branch_index": branch_index,
                            "fork_position": fork_position,
                            "trigger_token_id": next_token,
                            "injected_token_ids": injected_tokens,
                            "token_ids": branch_tokens,
                            "text": self.tokenizer.decode(branch_tokens, skip_special_tokens=False),
                        }
                    )
                    yield {
                        "type": "branch_done",
                        "branch_index": branch_index,
                        "total": len(branch_tokens),
                    }
                    yield {
                        "type": "scheduler",
                        "execution_mode": "sequential",
                        "phase": "main_resumed",
                        "main_active": True,
                        "active_branches": 0,
                        "queued_branches": 0,
                        "completed_branches": len(branches),
                        "batch_size": 1,
                    }

            if next_token == self.eos_token_id or len(main_tokens) >= self.max_new_tokens:
                break
            logical_position = prompt_length + response_position
            output = self.model(
                input_ids=torch.tensor([[next_token]], dtype=torch.long, device=self.device),
                position_ids=self._position_ids(
                    torch.tensor([[logical_position]], dtype=torch.long, device=self.device),
                    multimodal=multimodal,
                ),
                past_key_values=cache,
                use_cache=True,
            )
            cache = output.past_key_values
            next_token = self._next_token(output.logits, branch=False)

        yield {"type": "main_done", "total": len(main_tokens)}

        yield {
            "type": "done",
            "execution_mode": "sequential",
            "main": self.tokenizer.decode(main_tokens, skip_special_tokens=False),
            "main_token_ids": main_tokens,
            "prompt_tokens": prompt_length,
            "branches": branches,
            "branch_limit_reached": dropped_triggers > 0,
            "dropped_branch_triggers": dropped_triggers,
            "peak_batch_size": 1,
        }

    def generate(
        self,
        messages: list[dict[str, Any]],
        *,
        execution_mode: str | None = None,
    ) -> dict[str, Any]:
        """Consume the event stream and return its final result."""
        result: dict[str, Any] | None = None
        for event in self.stream(messages, execution_mode=execution_mode):
            if event.get("type") == "done":
                result = dict(event)
        if result is None:
            raise RuntimeError("Transformers decoding ended without a final result.")
        result.pop("type", None)
        return result


def _dtype(name: str) -> torch.dtype:
    return {
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
        "float32": torch.float32,
    }[name]


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run PaDoc Transformers inference.")
    parser.add_argument("--model", required=True)
    parser.add_argument("--query", default="Parse this document.")
    parser.add_argument("--image", action="append", default=[])
    parser.add_argument("--max-new-tokens", type=int, default=512)
    parser.add_argument("--max-branch-tokens", type=int, default=None)
    parser.add_argument("--max-concurrent-branches", type=int, default=8)
    parser.add_argument("--max-total-branches", type=int, default=64)
    parser.add_argument("--execution-mode", choices=EXECUTION_MODES, default="parallel")
    parser.add_argument(
        "--device",
        default="cuda:0",
        help="Single inference device (default: cuda:0); automatic sharding is not used.",
    )
    parser.add_argument(
        "--dtype",
        choices=("bfloat16", "float16", "float32"),
        default="bfloat16",
    )
    parser.add_argument(
        "--attn-implementation",
        choices=("sdpa", "eager", "flash_attention_2"),
        default="sdpa",
    )
    parser.add_argument("--json", action="store_true", dest="json_output")
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> None:
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    args = parse_args(argv)
    logger.info("Loading %s", args.model)
    engine = SequentialPaDocEngine.from_pretrained(
        args.model,
        dtype=_dtype(args.dtype),
        device=args.device,
        attn_implementation=args.attn_implementation,
        max_new_tokens=args.max_new_tokens,
        max_branch_tokens=args.max_branch_tokens,
        max_concurrent_branches=args.max_concurrent_branches,
        max_total_branches=args.max_total_branches,
        execution_mode=args.execution_mode,
    )
    content = [{"type": "image", "image": image} for image in args.image]
    content.append({"type": "text", "text": args.query})
    logger.info("Model devices: %s", ", ".join(engine.devices))
    result = engine.generate(
        [{"role": "user", "content": content}], execution_mode=args.execution_mode
    )
    if args.json_output:
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return
    print("MAIN")
    print(result["main"])
    for branch in result["branches"]:
        print(f"\nBRANCH {branch['branch_index']}")
        print(branch["text"])


if __name__ == "__main__":
    main()