from textual.app import ComposeResult, Widget
from textual.containers import Horizontal
from textual.widgets import Label
from textual.color import Color

from ai_hydra.constants.DHydraTui import DField
from ai_hydra.utils.HydraMetrics import HydraMetrics
from ai_hydra.utils.MetricEvent import MemEvent

BUCKET = "bucket"
NUM_BUCKETS = 20
BASE_COLOR = "#0f1e10"


def percent_to_hex(p: float) -> str:
    """
    Convert a float in [0.0, 1.0] to a hex string "00" - "ff".
    """
    p = max(0.0, min(1.0, p))
    value = int(round(p * 255))
    return f"{value:02x}"


class ATHMemory(Widget):

    def __init__(self, metrics: HydraMetrics, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.metrics = metrics

    def compose(self) -> ComposeResult:
        with Horizontal():
            yield Label("Memory", id=DField.REPLAY_MEM)
            for i in range(NUM_BUCKETS):
                yield Label(id=f"b{i}", classes=BUCKET)

    def refresh_data(self) -> None:

        mem_event = self.metrics.get_bucket_snaphot()
        if mem_event is None:
            return

        bucket_counts = mem_event.bucket_counts
        max_size = max(bucket_counts)
        min_size = min(bucket_counts)
        span = max_size - min_size

        bucket_idx = 0
        for count in bucket_counts:
            label_id = f"b{bucket_idx}"
            cur_label = self.query_one(f"#{label_id}", Label)
            cur_label.update(str(count))

            pct = 1.0 if span == 0 else (count - min_size) / span
            alpha_hex = percent_to_hex(pct)
            cur_label.styles.background = Color.parse(
                f"{BASE_COLOR}{alpha_hex}"
            )
            bucket_idx += 1
