# Copyright (c) 2025, Tri Dao.
from typing import Optional, Tuple, Dict, Any
from dataclasses import dataclass

import torch
from torch import Tensor

import cutlass.cute as cute
from cutlass import Int32
from cutlass.cute.runtime import from_dlpack, make_ptr

from quack.cute_dsl_utils import torch2cute_dtype_map
from quack.varlen_utils import VarlenArguments
from quack.tile_scheduler import TileSchedulerOptions


@dataclass
class GemmTensorInfo:
    tensor: Optional[Tensor]
    dtype: Optional[Any] = None
    major: Optional[str] = None
    cute_tensor: Optional[cute.Tensor] = None


class GemmWrapperBase:
    @staticmethod
    def validate_tensor(tensor: Tensor, name: str, ndim: int) -> None:
        assert tensor.dim() == ndim and tensor.is_cuda, f"{name} must be a {ndim}D CUDA tensor"
        assert tensor.dtype in torch2cute_dtype_map, f"Unsupported dtype for {name}"

    @staticmethod
    def validate_shape(tensor: Tensor, expected_shape: Tuple[int, ...], name: str) -> None:
        assert tensor.shape == expected_shape, (
            f"{name} must have shape {expected_shape}, got {tensor.shape}"
        )

    @staticmethod
    def get_major_order(tensor: Tensor, dims: Tuple[str, str, str]) -> str:
        # Tensor is already permuted to (dims[0], dims[1], dims[2])
        # stride(1) == 1 means dims[1] is contiguous (innermost)
        return dims[1] if tensor.stride(1) == 1 else dims[0]

    @staticmethod
    def create_cute_tensor(
        tensor: Optional[Tensor],
        major: Optional[str],
        dims: Tuple[str, str, str],
        assumed_align: int = 16,
    ) -> Optional[cute.Tensor]:
        if tensor is None:
            return None
        # Tensor is already permuted to (dims[0], dims[1], dims[2]) or (dim[0], dim[1])
        # If major is dims[1], leading_dim is 1; if major is dims[0], leading_dim is 0
        leading_dim = 1 if major == dims[1] else 0
        return from_dlpack(tensor.detach(), assumed_align=assumed_align).mark_layout_dynamic(
            leading_dim=leading_dim
        )

    @staticmethod
    def validate_and_prepare_tensors(
        A: Tensor,
        B: Tensor,
        D: Optional[Tensor] = None,
        C: Optional[Tensor] = None,
        additional_tensors: Optional[Dict[str, Tensor]] = None,
        cu_seqlens_m: Optional[Tensor] = None,
        cu_seqlens_k: Optional[Tensor] = None,
        A_idx: Optional[Tensor] = None,
    ) -> Tuple[int, int, int, int, Dict[str, GemmTensorInfo]]:
        assert not (cu_seqlens_m is not None and cu_seqlens_k is not None), (
            "Only one of cu_seqlens_m and cu_seqlens_k can be specified"
        )
        assert B.dtype == A.dtype, "A and B must have the same dtype"

        # Validate A_idx if provided (for gather_A case)
        gather_A = A_idx is not None
        if gather_A:
            assert cu_seqlens_m is not None or cu_seqlens_k is not None, (
                "gather_A requires either varlen_m or varlen_k"
            )
            assert A_idx.dtype == torch.int32, f"A_idx must be int32, got {A_idx.dtype}"
            assert A_idx.dim() == 1, f"A_idx must be 1D, got {A_idx.dim()}D"

        # Determine mode and extract dimensions
        if cu_seqlens_m is not None:
            # varlen_m: A is (total_m, k) or (whatever, k) if gather_A, B is (l, n, k), D/C are (total_m, n)
            assert A.dim() == 2, f"A must be 2D when using varlen_m, got {A.dim()}D"
            assert B.dim() == 3, f"B must be 3D with varlen_m, got {B.dim()}D"

            if gather_A:
                # When gather_A, A can have any number of rows, we use A_idx.shape[0] as total_M
                total_M = A_idx.shape[0]
                _, K = A.shape
            else:
                total_M, K = A.shape

            L, N, K_B = B.shape
            assert K == K_B, f"K dimension mismatch: A has {K}, B has {K_B}"
            assert cu_seqlens_m.shape == (L + 1,), (
                f"cu_seqlens_m must have shape ({L + 1},), got {cu_seqlens_m.shape}"
            )
            M = total_M
            dc_shape = (total_M, N)
            dc_ndim = 2
        elif cu_seqlens_k is not None:
            # varlen_k: A is (m, total_k) or (m, whatever) if gather_A, B is (n, total_k), D/C are (l, m, n)
            assert A.dim() == 2, f"A must be 2D when using varlen_k, got {A.dim()}D"
            assert B.dim() == 2, f"B must be 2D with varlen_k, got {B.dim()}D"

            if gather_A:
                # When gather_A with varlen_k, A can have any number of columns, we use A_idx.shape[0] as total_K
                M, _ = A.shape
                total_K = A_idx.shape[0]
            else:
                M, total_K = A.shape

            N, K_B = B.shape
            assert total_K == K_B, f"K dimension mismatch: expected {total_K}, B has {K_B}"
            L = cu_seqlens_k.shape[0] - 1
            assert cu_seqlens_k.shape == (L + 1,), (
                f"cu_seqlens_k must have shape ({L + 1},), got {cu_seqlens_k.shape}"
            )
            K = total_K
            dc_shape = (L, M, N)
            dc_ndim = 3
        else:
            # Normal case - all tensors must be 3D
            GemmWrapperBase.validate_tensor(A, "A", 3)
            GemmWrapperBase.validate_tensor(B, "B", 3)
            L, M, K = A.shape
            _, N, K_B = B.shape
            assert K == K_B, f"K dimension mismatch: A has {K}, B has {K_B}"
            GemmWrapperBase.validate_shape(B, (L, N, K), "B")
            dc_shape = (L, M, N)
            dc_ndim = 3

        # Validate D and C shapes uniformly
        for tensor, name in [(D, "D"), (C, "C")]:
            if tensor is not None:
                assert tensor.dim() == dc_ndim, (
                    f"{name} must be {dc_ndim}D for this mode, got {tensor.dim()}D"
                )
                assert tensor.shape == dc_shape, (
                    f"{name} shape {tensor.shape} doesn't match expected {dc_shape}"
                )

        tensors = {
            "A": GemmTensorInfo(A),
            "B": GemmTensorInfo(B),
            "D": GemmTensorInfo(D),
            "C": GemmTensorInfo(C),
        }

        if additional_tensors:
            for name, tensor in additional_tensors.items():
                if tensor is not None:
                    assert tensor.dim() == dc_ndim, (
                        f"{name} must be {dc_ndim}D for this mode, got {tensor.dim()}D"
                    )
                    assert tensor.shape == dc_shape, (
                        f"{name} shape {tensor.shape} doesn't match expected {dc_shape}"
                    )
                tensors[name] = GemmTensorInfo(tensor)

        return L, M, K, N, tensors

    @staticmethod
    def permute_tensors(
        tensors: Dict[str, GemmTensorInfo], varlen_m: bool = False, varlen_k: bool = False
    ) -> None:
        # Determine which tensors need permutation
        if varlen_m:
            # Only B needs permutation (3D tensor)
            tensors_to_permute = ["B"]
        elif varlen_k:
            # Only D and C need permutation (3D tensors)
            tensors_to_permute = ["D", "C"]
        else:
            # All tensors need permutation
            tensors_to_permute = None

        # Apply permutation from (L, *, *) -> (*, *, L) for selected tensors
        for name, info in tensors.items():
            if info.tensor is not None and info.tensor.ndim == 3:
                if tensors_to_permute is None or name in tensors_to_permute:
                    info.tensor = info.tensor.permute(1, 2, 0)

    @staticmethod
    def extract_dtypes(tensors: Dict[str, GemmTensorInfo]) -> None:
        for name, info in tensors.items():
            if info.tensor is not None:
                info.dtype = torch2cute_dtype_map[info.tensor.dtype]

    @staticmethod
    def determine_major_orders(
        tensors: Dict[str, GemmTensorInfo], major_configs: Dict[str, Tuple[str, str, str]]
    ) -> None:
        for name, dims in major_configs.items():
            if name in tensors and tensors[name].tensor is not None:
                tensors[name].major = GemmWrapperBase.get_major_order(tensors[name].tensor, dims)

    @staticmethod
    def create_cute_tensors(
        tensors: Dict[str, GemmTensorInfo], major_configs: Dict[str, Tuple[str, str, str]]
    ) -> None:
        for name, info in tensors.items():
            if info.tensor is not None and name in major_configs:
                info.cute_tensor = GemmWrapperBase.create_cute_tensor(
                    info.tensor, info.major, major_configs[name]
                )

    @staticmethod
    def create_scheduler_args(
        max_active_clusters: int,
        tile_count_semaphore: Optional[Tensor] = None,
        batch_idx_permute: Optional[Tensor] = None,
        max_swizzle_size: int = 8,
    ) -> TileSchedulerOptions:
        return TileSchedulerOptions(
            Int32(max_active_clusters),
            tile_count_semaphore=make_ptr(
                Int32, tile_count_semaphore.data_ptr(), cute.AddressSpace.gmem, assumed_align=4
            )
            if tile_count_semaphore is not None
            else None,
            batch_idx_permute=(
                from_dlpack(batch_idx_permute, assumed_align=4).mark_layout_dynamic(leading_dim=0)
            )
            if batch_idx_permute is not None
            else None,
            max_swizzle_size=Int32(max_swizzle_size),
        )

    @staticmethod
    def create_varlen_args(
        cu_seqlens_m: Optional[Tensor],
        cu_seqlens_k: Optional[Tensor],
        A_idx: Optional[Tensor],
    ) -> Optional[Any]:
        if cu_seqlens_m is None and cu_seqlens_k is None:
            return None
        return VarlenArguments(
            mCuSeqlensM=(
                from_dlpack(cu_seqlens_m, assumed_align=4).mark_layout_dynamic(leading_dim=0)
                if cu_seqlens_m is not None
                else None
            ),
            mCuSeqlensK=(
                from_dlpack(cu_seqlens_k, assumed_align=4).mark_layout_dynamic(leading_dim=0)
                if cu_seqlens_k is not None
                else None
            ),
            mAIdx=(
                from_dlpack(A_idx, assumed_align=4).mark_layout_dynamic(leading_dim=0)
                if A_idx is not None
                else None
            ),
        )

    @staticmethod
    def get_compile_key(
        tensors: Dict[str, GemmTensorInfo],
        activation: Optional[str],
        tile_shape_mn: Tuple[int, int],
        cluster_shape_mnk: Tuple[int, int, int],
        pingpong: bool,
        persistent: bool,
        has_semaphore: bool,
        *args,
        key_tensor_names: Tuple[str, ...] = ("A", "B", "D", "C"),
    ) -> Tuple:
        key_parts = []
        for name in key_tensor_names:
            if name in tensors:
                key_parts.append(tensors[name].dtype)
        key_parts.append(activation)
        key_parts.extend([tile_shape_mn, cluster_shape_mnk])
        for name in key_tensor_names:
            if name in tensors:
                key_parts.append(tensors[name].major)
        key_parts.extend([pingpong, persistent, has_semaphore])
        key_parts.extend(args)
        return tuple(key_parts)
