from io import BytesIO
from typing import Tuple

import numpy as np
from PIL import Image
import openslide
from skimage import color


def _level_downsample(slide: openslide.OpenSlide, level: int) -> float:
    return float(slide.level_downsamples[level])


def _level_dims(slide: openslide.OpenSlide, level: int) -> Tuple[int, int]:
    return slide.level_dimensions[level]


def _blank_tile(tile_size: int) -> bytes:
    img = Image.new("RGB", (tile_size, tile_size), (255, 255, 255))
    buf = BytesIO()
    img.save(buf, format="JPEG", quality=85)
    return buf.getvalue()


def _channel_from_rgb(rgb: Image.Image, channel: str) -> Image.Image:
    if channel == "original":
        return rgb

    arr = np.asarray(rgb).astype(np.float32) / 255.0
    hed = color.rgb2hed(arr)

    if channel == "hematoxylin":
        hed_only = np.zeros_like(hed)
        hed_only[..., 0] = hed[..., 0]
    elif channel == "eosin":
        hed_only = np.zeros_like(hed)
        hed_only[..., 1] = hed[..., 1]
    else:
        return rgb

    rgb_stain = color.hed2rgb(hed_only)
    rgb_stain = np.clip(rgb_stain, 0.0, 1.0)

    # Normalize for better contrast
    min_val = rgb_stain.min(axis=(0, 1), keepdims=True)
    max_val = rgb_stain.max(axis=(0, 1), keepdims=True)
    denom = np.maximum(max_val - min_val, 1e-6)
    rgb_stain = (rgb_stain - min_val) / denom

    out = (rgb_stain * 255.0).clip(0, 255).astype(np.uint8)
    return Image.fromarray(out, mode="RGB")


def get_tile_jpeg(
    slide: openslide.OpenSlide,
    level: int,
    x: int,
    y: int,
    tile_size: int,
    channel: str = "original",
) -> bytes:
    if level < 0 or level >= slide.level_count:
        raise ValueError("Invalid level")

    level_w, level_h = _level_dims(slide, level)

    px = x * tile_size
    py = y * tile_size

    if px >= level_w or py >= level_h:
        return _blank_tile(tile_size)

    downsample = _level_downsample(slide, level)
    x0 = int(px * downsample)
    y0 = int(py * downsample)

    region = slide.read_region((x0, y0), level, (tile_size, tile_size))
    rgb = region.convert("RGB")
    rgb = _channel_from_rgb(rgb, channel)

    buf = BytesIO()
    rgb.save(buf, format="JPEG", quality=85)
    return buf.getvalue()


def get_thumbnail_jpeg(
    slide: openslide.OpenSlide,
    size: int = 256,
    channel: str = "original",
) -> bytes:
    level = max(slide.level_count - 1, 0)
    level_w, level_h = _level_dims(slide, level)

    # Read the full lowest-resolution level, then downscale to a fixed thumbnail.
    region = slide.read_region((0, 0), level, (level_w, level_h))
    rgb = region.convert("RGB")
    rgb = _channel_from_rgb(rgb, channel)
    # Preserve aspect ratio and pad to square so the whole slide is visible.
    scale = min(size / max(level_w, 1), size / max(level_h, 1))
    new_w = max(int(level_w * scale), 1)
    new_h = max(int(level_h * scale), 1)
    rgb = rgb.resize((new_w, new_h), resample=Image.BILINEAR)

    canvas = Image.new("RGB", (size, size), (255, 255, 255))
    offset_x = (size - new_w) // 2
    offset_y = (size - new_h) // 2
    canvas.paste(rgb, (offset_x, offset_y))
    rgb = canvas

    buf = BytesIO()
    rgb.save(buf, format="JPEG", quality=85)
    return buf.getvalue()