import os
import hashlib
from typing import Optional


CUDA_HEADER = """\
#include <stdint.h>
#include <math.h>
#include <float.h>
#include <cuda_fp16.h>

__device__ float __sigmoid(float x) {
  return 1.0f / (1.0f + expf(-x));
}

"""

HEADER_GUARD_TMPL = """\
#ifndef {guard}
#define {guard}

{decls}

#endif
"""


class CudaSourceBuilder:
  def __init__(self):
    self.kernels: dict[str, str] = {}
    self.headers: list[str] = [CUDA_HEADER]
    self._decls: list[str] = []

  def add_kernel(self, name: str, source: str, launch_params: Optional[dict] = None):
    self.kernels[name] = source
    decl = self._make_decl(name, source)
    if decl:
      self._decls.append(decl)

  def _make_decl(self, name: str, source: str) -> str:
    for line in source.splitlines():
      if line.startswith("__global__"):
        sig = line.rstrip(" {")
        return sig + ";"
    return ""

  def build_source(self) -> str:
    parts = []
    for h in self.headers:
      parts.append(h)
    for src in self.kernels.values():
      parts.append(src)
      parts.append("\n")
    return "\n".join(parts)

  def build_header(self, guard_name: Optional[str] = None) -> str:
    if not guard_name:
      guard_name = "TRITONCU_GENERATED_H"
    decls = "\n".join(self._decls)
    return HEADER_GUARD_TMPL.format(guard=guard_name, decls=decls)

  def write_to_disk(
    self,
    output_dir: str,
    base_name: str = "kernels",
    write_header: bool = True,
  ):
    os.makedirs(output_dir, exist_ok=True)
    cu_path = os.path.join(output_dir, f"{base_name}.cu")
    h_path = os.path.join(output_dir, f"{base_name}.cuh")

    src = self.build_source()
    with open(cu_path, "w") as f:
      f.write(src)

    if write_header:
      guard = base_name.upper().replace("-", "_").replace("/", "_") + "_CUH"
      header = self.build_header(guard_name=guard)
      with open(h_path, "w") as f:
        f.write(header)

    return cu_path, h_path if write_header else None

  def source_hash(self) -> str:
    src = self.build_source()
    return hashlib.md5(src.encode()).hexdigest()[:8]
