from typing import Optional
from pydantic import BaseModel

from rich.live import Live
from rich.console import Console, Group
from rich.columns import Columns
from rich.panel import Panel

import numpy as np

class Cell (BaseModel):
    cell_value:str
    cell_color:Optional[str]=None
    rgb_array_color:tuple[int,int,int]

    debug_value:Optional[str]=None
    debug_color:Optional[str]=None

    def _render_rgb_array(self):
        return np.array(self.rgb_array_color, dtype=np.uint8),
    
    def _render_human(self, debug:bool):
        render_value = (
            self.debug_value if self.debug_value is not None else self.cell_value
        ) if debug else self.cell_value
        render_color = (
            self.debug_color if self.debug_color is not None else self.cell_color
        ) if debug else self.cell_color

        if render_color is None: return self._render_ansi(debug=debug)
        return f"[{render_color}]{render_value}[/{render_color}]"

    def _render_ansi(self, debug:bool):
        render_value = (
            self.debug_value if self.debug_value is not None else self.cell_value
        ) if debug else self.cell_value
        return render_value

    def render(self, mode="human", debug:bool=False):
        match mode:
            case "rgb_array": return self._render_rgb_array()
            case "human": return self._render_human(debug=debug)
            case "ansi": return self._render_ansi(debug=debug)
            case _: raise NotImplementedError(f"'{mode}' is not a valid mode.")

class Grid (BaseModel):
    size:tuple[int,int]=(5,5)
    panel_style:str="white"
    panel_style_debug:str="red"
    
    @property
    def height(self): return self.size[0]

    @property
    def width(self): return self.size[1]

    def _render_rgb_array(self, cells:list[list[Cell]]):
        return np.array([
            [cell.render(mode="rgb_array") for cell in row]
            for row in cells
        ], dtype=np.uint8)

    def _render_human(self, cells:list[list[Cell]], debug:bool):
        grid = [[
            cell.render(mode="human", debug=debug) for cell in row
        ] for row in cells]
        panel = Panel.fit(
            Group(*[
                Columns(row, equal=True) for row in grid #type:ignore
            ])
        )
        panel.style=self.panel_style_debug if debug else self.panel_style
        panel.height = self.height + 2
        panel.width = (self.width * 2) + 3
        return panel

    def _render_ansi(self, cells:list[list[Cell]], debug:bool):
        return "\n".join([ #type:ignore
            "".join([cell.render(mode="ansi", debug=debug) for cell in row]) #type:ignore
            for row in cells
        ])

    def _validate_input(self, cells:list[list[Cell]]):
        if len(cells) != self.height: return False
        for row in cells:
            if len(row) != self.width or not all(isinstance(cell, Cell) for cell in row):
                return False
        return True

    def render(self, cells:list[list[Cell]], mode:str="human", debug:bool=False):
        assert self._validate_input(cells), "grid.render:: invalid input."
        match mode:
            case "rgb_array": return self._render_rgb_array(cells)
            case "human": return self._render_human(cells, debug)
            case "ansi": return self._render_ansi(cells, debug)
            case _: raise NotImplementedError(f"'{mode}' is not a valid mode.")
