import io
import base64
import numpy as np
from PIL import Image

class TACAnalyzer:
    @staticmethod
    def generate_previews(cmyk_bytes: bytes, width: int, height: int, max_preview_dim: int = 800):
        raw_arr = np.frombuffer(cmyk_bytes, dtype=np.uint8).reshape((height, width, 4))
        scale = min(1.0, max_preview_dim / max(width, height))

        if scale < 1.0:
            step = max(1, int(1.0 / scale))
            arr = raw_arr[::step, ::step, :].copy()
            new_h, new_w, _ = arr.shape
        else:
            arr = raw_arr
            new_h, new_w = height, width

        c = arr[:, :, 0].astype(np.float32) / 255.0
        m = arr[:, :, 1].astype(np.float32) / 255.0
        y = arr[:, :, 2].astype(np.float32) / 255.0
        k = arr[:, :, 3].astype(np.float32) / 255.0

        r = (1.0 - c) * (1.0 - k) * 255.0
        g = (1.0 - m) * (1.0 - k) * 255.0
        b = (1.0 - y) * (1.0 - k) * 255.0
        rgb_arr = np.clip(np.dstack([r, g, b]), 0, 255).astype(np.uint8)
        img_composite = Image.fromarray(rgb_arr, "RGB")

        c_sub = (arr[:, :, 0].astype(np.float32) * 0.1).astype(np.uint8)
        cyan_rgb = np.dstack([255 - arr[:, :, 0], 255 - c_sub, np.full((new_h, new_w), 255, dtype=np.uint8)])
        m_sub = (arr[:, :, 1].astype(np.float32) * 0.2).astype(np.uint8)
        mag_rgb = np.dstack([np.full((new_h, new_w), 255, dtype=np.uint8), 255 - arr[:, :, 1], 255 - m_sub])
        yel_rgb = np.dstack([np.full((new_h, new_w), 255, dtype=np.uint8), np.full((new_h, new_w), 255, dtype=np.uint8), 255 - arr[:, :, 2]])
        blk_rgb = np.dstack([255 - arr[:, :, 3], 255 - arr[:, :, 3], 255 - arr[:, :, 3]])

        tac_pct = (c + m + y + k) * 100.0
        heatmap_rgb = np.zeros((new_h, new_w, 3), dtype=np.uint8)
        mask_safe = tac_pct <= 150.0
        heatmap_rgb[mask_safe] = [34, 197, 94]
        mask_mod = (tac_pct > 150.0) & (tac_pct <= 260.0)
        heatmap_rgb[mask_mod] = [59, 130, 246]
        mask_warn = (tac_pct > 260.0) & (tac_pct <= 320.0)
        heatmap_rgb[mask_warn] = [234, 179, 8]
        mask_crit = tac_pct > 320.0
        heatmap_rgb[mask_crit] = [239, 68, 68]

        def to_b64(img: Image.Image) -> str:
            buf = io.BytesIO()
            img.save(buf, format="JPEG", quality=85)
            return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode("ascii")

        return {
            "composite": to_b64(img_composite),
            "cyan": to_b64(Image.fromarray(cyan_rgb, "RGB")),
            "magenta": to_b64(Image.fromarray(mag_rgb, "RGB")),
            "yellow": to_b64(Image.fromarray(yel_rgb, "RGB")),
            "black": to_b64(Image.fromarray(blk_rgb, "RGB")),
            "tac_heatmap": to_b64(Image.fromarray(heatmap_rgb, "RGB"))
        }