"""Validated text → independent random glyphs → a white 24-bit BMP page."""
from __future__ import annotations

import io
import secrets
from pathlib import Path

import torch
from PIL import Image

from generator import Generator, LATENT_DIM

ALLOWED = frozenset('ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 \n')
MAX_TEXT = 4000
PAGE_SIZE = (1240, 1754)  # A4, approximately 150 DPI
MARGIN = 80


def validate_text(value: object) -> str:
    if not isinstance(value, str):
        raise ValueError('Enter text using A–Z, 0–9, spaces, and newlines.')
    # Normalize browser/Windows line endings; tabs and other whitespace stay invalid.
    value = value.replace('\r\n', '\n')
    if len(value) > MAX_TEXT:
        raise ValueError(f'Text is limited to {MAX_TEXT} characters.')
    invalid = sorted(set(value) - ALLOWED)
    if invalid:
        shown = ', '.join(repr(char) for char in invalid[:8])
        raise ValueError(f'Unsupported characters: {shown}. Use only A–Z, 0–9, spaces, and newlines.')
    if not any(char.isascii() and char.isalnum() for char in value):
        raise ValueError('Enter at least one uppercase letter or digit.')
    return value


def integer(value: object, name: str, minimum: int, maximum: int) -> int:
    if type(value) is not int or not minimum <= value <= maximum:
        raise ValueError(f'{name} must be an integer from {minimum} to {maximum}.')
    return value


def tile_layout(text: str, columns: int) -> list[str]:
    """Preserve every character: spaces occupy tiles, newlines start new rows."""
    rows = []
    for line in text.split('\n'):
        if not line:
            rows.append('')
        else:
            while len(line) > columns:
                # Prefer a word boundary; retain the separating space as a tile.
                boundary = line.rfind(' ', 0, columns)
                end = boundary + 1 if boundary >= 0 else columns
                rows.append(line[:end])
                line = line[end:]
            if line:
                rows.append(line)
    return rows


class LetterRenderer:
    def __init__(self, model_path: Path, device: str = 'auto'):
        torch.set_num_threads(2)
        selected = ('cuda' if torch.cuda.is_available() else 'cpu') if device == 'auto' else device
        if selected == 'cuda' and not torch.cuda.is_available():
            raise RuntimeError('CUDA was requested but is unavailable. Use --device cpu or auto.')
        self.device = torch.device(selected)
        checkpoint = torch.load(model_path, map_location='cpu', weights_only=True)
        if checkpoint['latent_dim'] != LATENT_DIM:
            raise ValueError('Model latent dimension must be 64.')
        self.classes = list(checkpoint['class_names'])
        if not set(ALLOWED - {' ', '\n'}).issubset(self.classes):
            raise ValueError('Model must include all uppercase letters and digits.')
        self.model = Generator(len(self.classes)).to(self.device).eval()
        self.model.load_state_dict(checkpoint['generator'])
        self.epoch = int(checkpoint['epoch'])
        self.weights = checkpoint['weights']

    def info(self) -> dict:
        return {'epoch': self.epoch, 'device': self.device.type, 'weights': self.weights,
                'latent_dim': LATENT_DIM, 'page_size': list(PAGE_SIZE), 'max_text': MAX_TEXT}

    @torch.inference_mode()
    def render(self, payload: dict) -> tuple[bytes, dict]:
        if not isinstance(payload, dict):
            raise ValueError('Request must be a JSON object.')
        text = validate_text(payload.get('text'))
        tile_size = integer(payload.get('tile_size', 28), 'Tile size', 28, 56)
        spacing = integer(payload.get('line_spacing', 12), 'Line spacing', 0, 32)
        requested_seed = payload.get('seed')
        seed = secrets.randbelow(2**32) if requested_seed is None else integer(requested_seed, 'Seed', 0, 2**32 - 1)
        width, height = PAGE_SIZE
        columns = (width - 2 * MARGIN) // tile_size
        rows = tile_layout(text, columns)
        capacity = (height - 2 * MARGIN + spacing) // (tile_size + spacing)
        if len(rows) > capacity:
            raise ValueError(f'Text needs {len(rows)} lines but this page fits {capacity}. Reduce tile size or line spacing, or shorten the text.')
        letters = [char for row in rows for char in row if char != ' ']
        class_ids = torch.tensor([self.classes.index(char) for char in letters], dtype=torch.long)
        # CPU RNG keeps the seed-to-latent mapping the same on CPU and CUDA.
        rng = torch.Generator(device='cpu').manual_seed(seed)
        noise = payload.get('_latents')
        if noise is None:
            noise = torch.randn(len(letters), LATENT_DIM, generator=rng)
        if noise.shape != (len(letters), LATENT_DIM):
            raise ValueError('Incorrect latent vector count.')
        # Preallocate the page; a space remains a completely white tile.
        page = Image.new('RGB', PAGE_SIZE, 'white')
        positions = [(MARGIN + col * tile_size, MARGIN + row_number * (tile_size + spacing))
                     for row_number, row in enumerate(rows) for col, char in enumerate(row) if char != ' ']
        for start in range(0, len(letters), 64):
            generated = self.model(noise[start:start + 64].to(self.device), class_ids[start:start + 64].to(self.device))
            # Invert native light-on-black output into dark ink on white paper.
            pixels = ((1 - generated.cpu()) * 127.5).clamp(0, 255).round().to(torch.uint8)
            for offset, glyph in enumerate(pixels):
                tile = Image.frombytes('L', (28, 28), bytes(glyph.flatten().tolist()))
                if tile_size != 28:
                    tile = tile.resize((tile_size, tile_size), Image.Resampling.NEAREST)
                page.paste(tile, positions[start + offset])
        result = io.BytesIO()
        page.save(result, format='BMP', dpi=(150, 150))
        return result.getvalue(), {'seed': seed, 'glyphs': len(letters), 'rows': len(rows),
                                   'columns': columns, 'epoch': self.epoch, 'width': width, 'height': height}