import collections
import math
import re
from typing import Any, Dict, Sequence

import torch
import triton
from torch.profiler import ProfilerActivity, profile

from .diff_engine import DiffCase


def _get_best_cuda_timing(timings_ms, key):
    """Look up the best CUDA-based timing for speedup calculation."""
    for provider in ("cuda", "compiled_cuda"):
        if provider in timings_ms and key in timings_ms[provider]:
            return timings_ms[provider][key]
    raise KeyError(f"No CUDA timing found for {key}")


def _shorten_kernel_name(name: str) -> str:
    """Strip template args and function params from CUDA kernel names.

    ``void motif::grouped_poly_norm_bwd_kernel<...>(...)``
    → ``motif::grouped_poly_norm_bwd_kernel``
    """
    # Remove leading 'void '
    s = re.sub(r"^void\s+", "", name)
    # Remove template args <...> (handles nested <>)
    while "<" in s:
        s = re.sub(r"<[^<>]*>", "", s)
    # Remove function params (...)
    s = re.sub(r"\(.*\)$", "", s)
    return s.strip()


def _compute_bytes(inputs, forward_fn, obj):
    """Compute total bytes: all input tensors read + all output tensors written."""
    input_bytes = sum(v.nbytes for v in inputs.values()
                      if isinstance(v, torch.Tensor))
    output = forward_fn()
    if isinstance(output, torch.Tensor):
        output_bytes = output.nbytes
    elif isinstance(output, (tuple, list)):
        output_bytes = sum(o.nbytes for o in output
                           if isinstance(o, torch.Tensor))
    else:
        output_bytes = 0
    return input_bytes + output_bytes


def profile_bench(fn, warmup=5, repeat=10, verbose=True, total_bytes=0):
    """Measure CUDA kernel time via torch.profiler.

    Profiles the function, sums all CUDA kernel durations, and returns
    the median across repeats.  Also prints a per-kernel breakdown when
    *verbose* is True so the caller can spot unexpected kernels.

    Parameters
    ----------
    total_bytes : int
        Total bytes transferred (inputs read + outputs written).
        If > 0, prints bandwidth in GB/s after the breakdown.

    Returns
    -------
    median_ms : float
        Median total CUDA kernel time in **milliseconds** (same unit as
        ``triton.testing.do_bench``).
    """
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()

    kernel_times_us: list[float] = []
    last_breakdown: list[tuple[str, float]] = []

    for _ in range(repeat):
        with profile(activities=[ProfilerActivity.CUDA]) as prof:
            fn()

        breakdown: dict[str, float] = {}
        for evt in prof.key_averages():
            if evt.device_time_total > 0:
                breakdown[evt.key] = (breakdown.get(evt.key, 0) +
                                      evt.device_time_total)

        total_us = sum(breakdown.values())
        kernel_times_us.append(total_us)
        last_breakdown = sorted(breakdown.items(),
                                key=lambda x: x[1],
                                reverse=True)

    median_us = sorted(kernel_times_us)[len(kernel_times_us) // 2]

    if verbose and last_breakdown:
        total = sum(t for _, t in last_breakdown)
        names = [_shorten_kernel_name(n) for n, _ in last_breakdown]
        col_w = max(len(n) for n in names) + 2
        col_w = max(col_w, len("Total kernel time") + 2)
        for name, (_, t) in zip(names, last_breakdown):
            pct = 100 * t / total if total > 0 else 0
            print(f"    {name:<{col_w}s} {t:>8.1f}us ({pct:4.1f}%)")
        print(f"    {'Total kernel time':<{col_w}s} {total:>8.1f}us")
        if total_bytes > 0 and median_us > 0:
            bw_gbs = total_bytes / (median_us * 1e-6) / 1e9
            print(f"    {'Bandwidth':<{col_w}s} {bw_gbs:>7.1f} GB/s"
                  f"  ({total_bytes / 1e6:.1f} MB)")

    return median_us / 1000  # us -> ms


def make_fwd_key(batch_size, seq_len, dim):
    return f"forward : ({batch_size}, {seq_len}, {dim})"


def make_bwd_key(batch_size, seq_len, dim):
    return f"backward : ({batch_size}, {seq_len}, {dim})"


def parse_config_string(config_str):
    match = re.match(r"(\w+)\s*:\s*\(\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\)",
                     config_str)
    if not match:
        raise ValueError(f"Invalid config string: {config_str}")
    _, bs, sl, d = match.groups()
    return int(bs), int(sl), int(d)


def make_fwd_benchmark_for_case(
    *,
    case: DiffCase,
    configs: Sequence[tuple[int, int, int]],
    plot_name: str,
    ylabel: str = "",
    line_vals=("naive", "cuda", "speedup"),
    line_names: Dict[str, str] | None = None,
    dtype=torch.bfloat16,
    eps: float = 1e-6,
    time_unit_scale: float = 1000,
):
    timings_ms = collections.defaultdict(dict)
    bytes_map: dict[str, int] = {}
    line_vals = list(line_vals)
    line_names = line_names or {v: v.title() for v in line_vals}
    x_vals = [list(_) for _ in configs]

    @triton.testing.perf_report(
        triton.testing.Benchmark(x_names=["dim", "batch_size", "seq_len"],
                                 x_vals=x_vals,
                                 line_arg="provider",
                                 line_vals=line_vals,
                                 line_names=[line_names[v] for v in line_vals],
                                 ylabel=ylabel,
                                 plot_name=plot_name,
                                 args={}))
    def bench(dim, batch_size, seq_len, provider):
        key = make_fwd_key(dim, batch_size, seq_len)
        I = case.build_inputs(batch_size, seq_len, dim, dtype, eps)
        if provider == "speedup":
            return round(
                timings_ms["naive"][key] /
                _get_best_cuda_timing(timings_ms, key), 2)
        if provider.endswith("_bw"):
            base = provider[:-3]
            ms = timings_ms[base][key]
            return round(bytes_map[key] / (ms * 1e-3) / 1e9, 2)
        if provider == "naive":
            obj = case.make_naive(I)
        elif provider == "compiled" and hasattr(case, "make_compiled"):
            obj = case.make_compiled(I)
        else:
            obj = case.make_cuda(I)
        run = lambda: case.forward(obj, I)
        nbytes = _compute_bytes(I, run, obj)
        bytes_map[key] = nbytes
        print(f"  [{provider}] {key}")
        ms = profile_bench(run, total_bytes=nbytes)
        timings_ms[provider][key] = ms
        return time_unit_scale * ms

    return bench


def make_fwd_benchmark_plot_for_case(
    *,
    case: DiffCase,
    configs: Sequence[tuple[int, int, int]],
    plot_name: str,
    ylabel: str = "Relative Speedup",
    line_vals=("naive", "cuda"),
    line_names: Dict[str, str] | None = None,
    dtype=torch.bfloat16,
    eps: float = 1e-6,
):
    timings_ms = collections.defaultdict(dict)
    spdup_ratio = list()
    line_vals = list(line_vals)
    line_names = line_names or {v: v.title() for v in line_vals}
    x_vals = [make_fwd_key(*_) for _ in configs]
    x_vals.append("Geometric Mean")

    @triton.testing.perf_report(
        triton.testing.Benchmark(x_names=["config"],
                                 x_vals=x_vals,
                                 line_arg="provider",
                                 line_vals=line_vals,
                                 line_names=[line_names[v] for v in line_vals],
                                 ylabel=ylabel,
                                 plot_name=plot_name,
                                 args={}))
    def bench(config, provider):
        if config == "Geometric Mean":
            if provider == "cuda":
                return round(math.prod(spdup_ratio)**(1 / len(spdup_ratio)), 2)
            else:
                return 1.00
        batch_size, seq_len, dim = parse_config_string(config)
        I = case.build_inputs(batch_size, seq_len, dim, dtype, eps)
        if provider == "naive":
            obj = case.make_naive(I)
        elif provider == "compiled" and hasattr(case, "make_compiled"):
            obj = case.make_compiled(I)
        else:
            obj = case.make_cuda(I)
        run = lambda: case.forward(obj, I)
        nbytes = _compute_bytes(I, run, obj)
        print(f"  [{provider}] {config}")
        ms = profile_bench(run, total_bytes=nbytes)
        timings_ms[provider][config] = ms
        if provider == "cuda":
            ratio = timings_ms["naive"][config] / _get_best_cuda_timing(
                timings_ms, config)
            spdup_ratio.append(ratio)
            return round(ratio, 2)
        else:
            return 1.00

    return bench


def make_bwd_benchmark_for_case(
    *,
    case: DiffCase,
    configs: Sequence[tuple[int, int, int]],
    plot_name: str,
    ylabel: str = "",
    line_vals=("naive", "cuda", "speedup"),
    line_names: Dict[str, str] | None = None,
    dtype=torch.bfloat16,
    eps: float = 1e-6,
    time_unit_scale: float = 1000,
):
    timings_ms = collections.defaultdict(dict)
    bytes_map: dict[str, int] = {}
    line_vals = list(line_vals)
    line_names = line_names or {v: v.title() for v in line_vals}
    x_vals = [list(_) for _ in configs]

    @triton.testing.perf_report(
        triton.testing.Benchmark(x_names=["dim", "batch_size", "seq_len"],
                                 x_vals=x_vals,
                                 line_arg="provider",
                                 line_vals=line_vals,
                                 line_names=[line_names[v] for v in line_vals],
                                 ylabel=ylabel,
                                 plot_name=plot_name,
                                 args={}))
    def bench(dim, batch_size, seq_len, provider):
        key = make_bwd_key(dim, batch_size, seq_len)
        I = case.build_inputs(batch_size, seq_len, dim, dtype, eps)
        if provider == "speedup":
            return round(
                timings_ms["naive"][key] /
                _get_best_cuda_timing(timings_ms, key), 2)
        if provider.endswith("_bw"):
            base = provider[:-3]
            ms = timings_ms[base][key]
            return round(bytes_map[key] / (ms * 1e-3) / 1e9, 2)
        if provider == "naive":
            obj = case.make_naive(I)
        elif provider == "compiled" and hasattr(case, "make_compiled"):
            obj = case.make_compiled(I)
        else:
            obj = case.make_cuda(I)
        y = case.forward(obj, I)
        gin = list(case.grad_inputs(I)) + list(obj.parameters())
        if isinstance(y, torch.Tensor):
            g = [torch.randn_like(y)]
        else:
            g = [torch.randn_like(r) for r in y]
        run = lambda: torch.autograd.grad(y,
                                          gin,
                                          g,
                                          retain_graph=True,
                                          create_graph=False,
                                          allow_unused=False)
        fwd_run = lambda: case.forward(obj, I)
        nbytes = _compute_bytes(I, fwd_run, obj)
        bytes_map[key] = nbytes
        print(f"  [{provider}] {key}")
        ms = profile_bench(run, total_bytes=nbytes)
        timings_ms[provider][key] = ms
        return time_unit_scale * ms

    return bench


def make_bwd_benchmark_plot_for_case(
    *,
    case: DiffCase,
    configs: Sequence[tuple[int, int, int]],
    plot_name: str,
    ylabel: str = "Relative Speedup",
    line_vals=("naive", "cuda"),
    line_names: Dict[str, str] | None = None,
    dtype=torch.bfloat16,
    eps: float = 1e-6,
):
    timings_ms = collections.defaultdict(dict)
    spdup_ratio = list()
    line_vals = list(line_vals)
    line_names = line_names or {v: v.title() for v in line_vals}
    x_vals = [make_bwd_key(*_) for _ in configs]
    x_vals.append("Geometric Mean")

    @triton.testing.perf_report(
        triton.testing.Benchmark(x_names=["config"],
                                 x_vals=x_vals,
                                 line_arg="provider",
                                 line_vals=line_vals,
                                 line_names=[line_names[v] for v in line_vals],
                                 ylabel=ylabel,
                                 plot_name=plot_name,
                                 args={}))
    def bench(config, provider):
        if config == "Geometric Mean":
            if provider == "cuda":
                return round(math.prod(spdup_ratio)**(1 / len(spdup_ratio)), 2)
            else:
                return 1.00
        batch_size, seq_len, dim = parse_config_string(config)
        I = case.build_inputs(batch_size, seq_len, dim, dtype, eps)
        if provider == "naive":
            obj = case.make_naive(I)
        elif provider == "compiled" and hasattr(case, "make_compiled"):
            obj = case.make_compiled(I)
        else:
            obj = case.make_cuda(I)
        y = case.forward(obj, I)
        gin = list(case.grad_inputs(I)) + list(obj.parameters())
        if isinstance(y, torch.Tensor):
            g = [torch.randn_like(y)]
        else:
            g = [torch.randn_like(r) for r in y]
        run = lambda: torch.autograd.grad(y,
                                          gin,
                                          g,
                                          retain_graph=True,
                                          create_graph=False,
                                          allow_unused=False)
        fwd_run = lambda: case.forward(obj, I)
        nbytes = _compute_bytes(I, fwd_run, obj)
        print(f"  [{provider}] {config}")
        ms = profile_bench(run, total_bytes=nbytes)
        timings_ms[provider][config] = ms
        if provider == "cuda":
            ratio = timings_ms["naive"][config] / _get_best_cuda_timing(
                timings_ms, config)
            spdup_ratio.append(ratio)
            return round(ratio, 2)
        else:
            return 1.00

    return bench