import json
from abc import ABC, abstractmethod
from pathlib import Path

import ml_dtypes  # noqa: F401  # register bfloat16 dtype with numpy
import numpy as np
import torch
import torch.distributed as dist

from bergson.utils.utils import convert_dtype_to_np, tensor_to_numpy


class ScoreWriter(ABC):
    """
    Base class for score writers.
    """

    @abstractmethod
    def __call__(
        self,
        indices: list[int],
        scores: torch.Tensor,
    ):
        """
        Write the scores to the score writer.
        """
        raise NotImplementedError("Subclasses must implement this method")

    @abstractmethod
    def flush(self):
        """
        Flush the score writer.
        """
        raise NotImplementedError("Subclasses must implement this method")


class InMemoryScoreWriter(ScoreWriter):
    """Stores scores in memory as a torch tensor."""

    def __init__(
        self, num_items: int, num_scores: int, dtype: torch.dtype = torch.float32
    ):
        self.scores = torch.zeros((num_items, num_scores), device="cpu", dtype=dtype)

    def __call__(self, indices: list[int], scores: torch.Tensor):
        self.scores[indices] = scores.to(dtype=self.scores.dtype).cpu()

    def flush(self):
        # No-op for in-memory storage
        pass


class MemmapScoreWriter(ScoreWriter):
    """
    Writes scores to a memory-mapped file on disk.

    Supports bfloat16 via ml_dtypes.
    """

    def __init__(
        self,
        path: Path,
        num_items: int,
        num_scores: int,
        *,
        dtype: torch.dtype = torch.float32,
        flush_interval: int = 64,
    ):
        self.path = path
        self.num_scores = num_scores
        self.dtype = dtype
        self.flush_interval = flush_interval
        self.num_batches_since_flush = 0

        self.path.mkdir(parents=True, exist_ok=True)
        scores_file_path = self.path / "scores.bin"

        # Convert torch dtype to numpy dtype (handles bfloat16 via ml_dtypes)
        np_dtype = convert_dtype_to_np(dtype)
        score_size = np_dtype.itemsize
        bool_size = np.dtype("bool").itemsize

        # Build a structured dtype with (score, written) pairs per query
        # Align each pair to the next power of 2 for efficiency
        pair_size = score_size + bool_size
        aligned_pair_size = 1 << (pair_size - 1).bit_length()  # Next power of 2

        names = []
        formats = []
        offsets = []
        for i in range(num_scores):
            names.append(f"score_{i}")
            formats.append(np_dtype)
            offsets.append(i * aligned_pair_size)

            names.append(f"written_{i}")
            formats.append("bool")
            offsets.append(i * aligned_pair_size + score_size)

        total_bytes = num_scores * aligned_pair_size
        # Round up to the nearest 8 bytes
        itemsize = ((total_bytes + 7) // 8) * 8

        # For JSON serialization, convert numpy dtype to string
        format_strs = [str(f) if isinstance(f, np.dtype) else f for f in formats]
        struct_dtype_json = {
            "names": names,
            "formats": format_strs,
            "offsets": offsets,
            "itemsize": itemsize,
        }

        struct_dtype = {
            "names": names,
            "formats": formats,
            "offsets": offsets,
            "itemsize": itemsize,
        }

        rank = dist.get_rank() if dist.is_initialized() else 0
        if rank == 0 and not scores_file_path.exists():
            print(f"Creating new scores file: {scores_file_path}")

            self.scores = np.memmap(
                str(scores_file_path),
                dtype=np.dtype(struct_dtype),  # type: ignore
                mode="w+",
                shape=(num_items,),
            )

            for name in names:
                if "written" in name:
                    self.scores[name][:] = False
            self.flush()

            # Persist metadata for future runs
            with (path / "info.json").open("w") as f:
                json.dump(
                    {
                        "num_items": num_items,
                        "num_scores": num_scores,
                        "dtype": struct_dtype_json,
                    },
                    f,
                    indent=2,
                )

        if dist.is_initialized():
            dist.barrier()

        self.scores = np.memmap(
            str(scores_file_path),
            dtype=np.dtype(struct_dtype),  # type: ignore
            mode="r+",
            shape=(num_items,),
        )

    def __call__(self, indices: list[int], scores: torch.Tensor):
        # scores: [num_indices, num_scores]
        scores = scores.to(dtype=self.dtype)
        for i in range(self.num_scores):
            score_col = tensor_to_numpy(scores[:, i].cpu()).flatten()
            self.scores[f"score_{i}"][indices] = score_col
            self.scores[f"written_{i}"][indices] = True

        self.num_batches_since_flush += 1
        if self.num_batches_since_flush >= self.flush_interval:
            self.flush()

    def flush(self):
        self.scores.flush()
        self.num_batches_since_flush = 0
