#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""NPU-specific encoder ACL graph: params, runtime context, FIA replay updates, and manager."""

from __future__ import annotations

from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any

import torch
import torch_npu
from vllm.logger import logger
from vllm.platforms import current_platform
from vllm.v1.worker.encoder_cudagraph import BudgetGraphMetadata, EncoderCudaGraphManager

from vllm_ascend.utils import vllm_version_is, weak_ref_tensors

# ---------------------------------------------------------------------------
# Per–encoder-budget ACL graph bookkeeping (ViT FIA tasks)
# ---------------------------------------------------------------------------


@dataclass
class EncoderGraphParams:
    """Mirrors :class:`vllm_ascend.compilation.acl_graph.GraphParams` but keyed by encoder token budget."""

    # TODO: Fully support upstream dual-path encoder graph on Ascend. The
    # current FIA bookkeeping is keyed only by token_budget; dual-path models
    # need graph params separated by (path, token_budget).
    events: dict[int, list[torch.npu.ExternalEvent]] = field(default_factory=dict)
    workspaces: dict[int, torch.Tensor | None] = field(default_factory=dict)
    handles: dict[int, list[Any]] = field(default_factory=dict)
    # Flattened per-forward insertion order (one entry per ViT block invocation).
    attn_params: dict[int, list[tuple]] = field(default_factory=dict)


_encoder_graph_params: EncoderGraphParams | None = None


def set_encoder_graph_params(token_budgets: list[int]) -> None:
    global _encoder_graph_params
    budgets_sorted_unique = sorted(token_budgets)
    _encoder_graph_params = EncoderGraphParams(
        events={b: [] for b in budgets_sorted_unique},
        workspaces={b: None for b in budgets_sorted_unique},
        handles={b: [] for b in budgets_sorted_unique},
        attn_params={b: [] for b in budgets_sorted_unique},
    )


def get_encoder_graph_params() -> EncoderGraphParams | None:
    return _encoder_graph_params


def update_encoder_graph_workspace(token_budget: int, workspace: torch.Tensor) -> None:
    if _encoder_graph_params is None:
        return
    _encoder_graph_params.workspaces[token_budget] = workspace


# ---------------------------------------------------------------------------
# Capture / replay runtime state (thread-local module singleton)
# ---------------------------------------------------------------------------


@dataclass
class EncoderForwardContext:
    """Vision encoder NPUGraph runtime flags and host-side FIA arguments.

    Captured tensors stay on device; FIA ``graph_task_update`` needs Python ``list[int]``
    lengths that are refreshed each replay from encoder metadata buffers on device (see RFC).
    """

    token_budget: int | None = None
    capturing: bool = False
    cu_seqlens_cpu: torch.Tensor | None = None


_context = EncoderForwardContext()


def get_encoder_forward_context() -> EncoderForwardContext:
    return _context


def _reset_encoder_forward_context() -> None:
    """Clear replay-time host length fields."""

    _context.token_budget = None
    _context.capturing = False
    _context.cu_seqlens_cpu = None


@contextmanager
def set_encoder_forward_context(
    token_budget: int,
    capturing: bool,
    *,
    cu_seqlens_cpu: torch.Tensor | None = None,
):
    """Enter encoder graph replay (FIA host args): callers must pass lengths each time.

    On exit, replay host fields are **cleared** (not restored). Tensors must not be reused
    across replays without repopulating from the current batch buffers.
    """

    _context.token_budget = token_budget
    _context.capturing = capturing
    _context.cu_seqlens_cpu = cu_seqlens_cpu
    try:
        yield _context
    finally:
        _reset_encoder_forward_context()


# ---------------------------------------------------------------------------
# FIA actual_seq_lengths (cu_seqlens -> list[int])
# ---------------------------------------------------------------------------


def maybe_compute_actual_seq_lengths(
    cu_seqlens: torch.Tensor,
    num_query_tokens: int,
    num_kv_tokens: int,
    *,
    cudagraph_mm_encoder: bool = False,
) -> tuple[list[int], list[int]]:
    """Convert ``cu_seqlens`` to FIA host ``actual_seq_lengths``.

    Drops the leading-zero marker; with ``cudagraph_mm_encoder`` filters endpoints and
    aligns the terminal to ``num_query_tokens``; when Q≠KV, scales Q endpoints by
    ``num_kv_tokens // num_query_tokens`` for the KV list.
    """
    flat = cu_seqlens.detach().cpu().view(-1).tolist()
    actual = flat[1:] if flat else flat

    if not cudagraph_mm_encoder:
        pass
    elif num_query_tokens <= 0:
        actual = [0]
    else:
        filtered: list[int] = []
        for end in actual:
            if end <= 0:
                continue
            if end > num_query_tokens:
                break
            if not filtered or end > filtered[-1]:
                filtered.append(end)

        if not filtered or filtered[-1] != num_query_tokens:
            filtered.append(num_query_tokens)
        actual = filtered

    if num_kv_tokens == num_query_tokens:
        return actual, actual

    assert num_kv_tokens % num_query_tokens == 0
    ratio = num_kv_tokens // num_query_tokens
    return actual, [end * ratio for end in actual]


def update_encoder_graph_params(
    update_stream: torch.npu.Stream,
    token_budget: int,
) -> None:
    """Re-bind fused infer attention host tensors inside the encoder NPUGraph (parallel to LLM path).

    This deliberately bypasses :class:`AttentionBackend` — ViT attention is not registered there — but reuses
    the same ``graph_task_update_{begin,end}`` + ``ExternalEvent`` ordering pattern as
    :meth:`AscendAttentionBackendImpl.update_graph_params`.
    """

    params = get_encoder_graph_params()
    if params is None or token_budget not in params.handles:
        return

    handles = params.handles[token_budget]
    events = params.events[token_budget]
    attn_blocks = params.attn_params[token_budget]
    workspace = params.workspaces.get(token_budget)

    if len(handles) != len(events) or len(handles) != len(attn_blocks):
        raise RuntimeError(
            "Encoder graph bookkeeping is inconsistent: "
            f"budget={token_budget} handles={len(handles)} "
            f"events={len(events)} attn_blocks={len(attn_blocks)}"
        )

    with torch.npu.stream(update_stream):
        for handle, event, packed in zip(handles, events, attn_blocks):
            (
                query,
                key,
                value,
                block_table,
                attn_mask,
                block_size,
                num_kv_heads,
                num_heads,
                scale,
                output,
                softmax_lse,
            ) = packed

            num_query_tokens = query.shape[0]
            num_kv_tokens = key.shape[0]
            cu_seqlens_cpu = get_encoder_forward_context().cu_seqlens_cpu

            actual_seq_lengths_q, actual_seq_lengths_kv = maybe_compute_actual_seq_lengths(
                cu_seqlens_cpu,
                num_query_tokens,
                num_kv_tokens,
                cudagraph_mm_encoder=True,
            )

            torch.npu.graph_task_update_begin(update_stream, handle)
            torch_npu.npu_fused_infer_attention_score.out(
                query=query,
                key=key,
                value=value,
                atten_mask=attn_mask,
                block_table=block_table,
                input_layout="TND",
                block_size=block_size,
                actual_seq_lengths=actual_seq_lengths_q,
                actual_seq_lengths_kv=actual_seq_lengths_kv,
                num_key_value_heads=num_kv_heads,
                num_heads=num_heads,
                scale=scale,
                sparse_mode=0,
                workspace=workspace,
                out=[output, softmax_lse],
            )
            torch.npu.graph_task_update_end(update_stream)
            event.record(update_stream)


# ---------------------------------------------------------------------------
# Encoder NPUGraph manager
# ---------------------------------------------------------------------------
class EncoderAclGraphManager(EncoderCudaGraphManager):
    """Hooks encoder capture/replay into Ascend FIA graph-task infrastructure."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.graph_pool = current_platform.get_global_graph_pool()
        self.update_stream: torch.npu.Stream | None = None

    def capture(self, graph_pool: Any | None = None):
        encoder_graph_pool = graph_pool if graph_pool is not None else self.graph_pool
        self.graph_pool = encoder_graph_pool

        set_encoder_graph_params(self.token_budgets)

        super().capture(graph_pool=encoder_graph_pool)

        weak_ref_workspaces()

    def _capture_budget_graph(self, token_budget: int, path: str = "default"):
        logger.debug(
            "Capturing encoder aclgraph for budget=%d, max_batch_size=%d, max_frames_per_batch=%d",
            token_budget,
            self.max_batch_size,
            self.max_frames_per_batch,
        )

        if vllm_version_is("0.23.0"):
            capture_inputs = self.model.prepare_encoder_cudagraph_capture_inputs(
                token_budget,
                self.max_batch_size,
                self.max_frames_per_batch,
                self.device,
                self.dtype,
            )
        else:
            capture_inputs = self.model.prepare_encoder_cudagraph_capture_inputs(
                token_budget,
                self.max_batch_size,
                self.max_frames_per_batch,
                self.device,
                self.dtype,
                path,
            )

        values = capture_inputs.values
        with torch.inference_mode():
            if vllm_version_is("0.23.0"):
                output = self.model.encoder_cudagraph_forward(dict(values))
            else:
                output = self.model.encoder_cudagraph_forward(dict(values), path=path)
            output_buffer = torch.empty_like(output)

        graph = torch.npu.NPUGraph()
        with (
            set_encoder_forward_context(token_budget, True),
            torch.inference_mode(),
            torch.npu.graph(graph, self.graph_pool),
        ):
            if vllm_version_is("0.23.0"):
                output = self.model.encoder_cudagraph_forward(dict(values))
            else:
                output = self.model.encoder_cudagraph_forward(dict(values), path=path)
            output_buffer.copy_(output)

        graph_meta = BudgetGraphMetadata(
            token_budget=token_budget,
            max_batch_size=self.max_batch_size,
            max_frames_per_batch=self.max_frames_per_batch,
            graph=graph,
            input_buffers=values,
            output_buffer=weak_ref_tensors(output_buffer),
        )
        if vllm_version_is("0.23.0"):
            self.budget_graphs[token_budget] = graph_meta
        else:
            graph_set = self._get_graph_set(path)
            graph_set[token_budget] = graph_meta

    def _run_budget_graph(
        self,
        mm_kwargs: dict[str, Any],
        token_budget: int,
        path: str = "default",
    ) -> torch.Tensor | None:
        num_items = len(self._get_item_specs(mm_kwargs))
        if vllm_version_is("0.23.0"):
            if token_budget not in self.budget_graphs:
                self.graph_misses += num_items
                return None
            graph_meta = self.budget_graphs[token_budget]
        else:
            graph_set = self._get_graph_set(path)
            if token_budget not in graph_set:
                self.graph_misses += num_items
                return None
            graph_meta = graph_set[token_budget]

        if vllm_version_is("0.23.0"):
            replay = self.model.prepare_encoder_cudagraph_replay_buffers(
                mm_kwargs,
                self.max_batch_size,
                self.max_frames_per_batch,
            )
            buffer_items = ((key, graph_meta.input_buffers[key]) for key in self.config.buffer_keys)
        else:
            replay = self.model.prepare_encoder_cudagraph_replay_buffers(
                mm_kwargs,
                self.max_batch_size,
                self.max_frames_per_batch,
                path,
            )
            buffer_items = graph_meta.input_buffers.items()

        for key, buf in buffer_items:
            src = replay.values.get(key)
            if src is None:
                continue
            if src.ndim == 0:
                buf.copy_(src)
            else:
                padding_logic = self.config.padding_logics.get(key, self._copy_padded_buffer)
                padding_logic(buf, src)

        cu_seqlens = graph_meta.input_buffers.get("cu_seqlens")
        cu_seqlens_cpu = None if cu_seqlens is None else cu_seqlens.cpu()

        update_stream = self.update_stream
        if update_stream is None:
            update_stream = torch.npu.Stream()

        graph_meta.graph.replay()

        with set_encoder_forward_context(
            token_budget,
            False,
            cu_seqlens_cpu=cu_seqlens_cpu,
        ):
            update_encoder_graph_params(update_stream, token_budget)

        self.graph_hits += num_items
        return graph_meta.output_buffer


def weak_ref_workspaces() -> None:
    params = get_encoder_graph_params()
    if params is None:
        return
    for budget, ws in list(params.workspaces.items()):
        if ws is None:
            continue
        params.workspaces[budget] = weak_ref_tensors(ws)