from __future__ import annotations
import math
from numbers import Integral
import torch
from torch import nn
GEN_IMAGE_ROW_MARKERS = 'gen_image_row_markers'
GEN_IMAGE_FIXED_PREFIX_TOKENS = 'gen_image_fixed_prefix_tokens'
GENERATION_SKIP_DOWNSAMPLE = 'generation_skip_down_sample'
ROW_MARKER_OFFSET = 0
IMAGE_MASK_OFFSET = 1
IMAGE_END_OFFSET = 2
FIXED_PREFIX_OFFSET = 3
RESERVED_GENERATION_ROWS = 256
MARKERS_PER_RAW_ROW = 2
UNCOMPRESSED_MARKERS_PER_RAW_ROW = 1
_SUPPORTED_MARKERS_PER_ROW = (UNCOMPRESSED_MARKERS_PER_RAW_ROW, MARKERS_PER_RAW_ROW)

def _validate_markers_per_row(value) -> int:
    if isinstance(value, bool) or not isinstance(value, Integral):
        raise TypeError('markers_per_row must be an integer')
    value = int(value)
    if value not in _SUPPORTED_MARKERS_PER_ROW:
        raise ValueError('markers_per_row must be 1 for an uncompressed image grid or 2 for a grid that will undergo 2x spatial compression')
    return value

def _validate_fixed_prefix_count(value) -> int:
    if value is None:
        return 0
    if isinstance(value, bool) or not isinstance(value, Integral):
        raise TypeError(f'config.{GEN_IMAGE_FIXED_PREFIX_TOKENS} must be an integer')
    value = int(value)
    max_tokens = RESERVED_GENERATION_ROWS - FIXED_PREFIX_OFFSET
    if not 0 <= value <= max_tokens:
        raise ValueError(f'config.{GEN_IMAGE_FIXED_PREFIX_TOKENS} must be between 0 and {max_tokens}')
    return value

def image_row_markers_enabled(config) -> bool:
    value = getattr(config, GEN_IMAGE_ROW_MARKERS, False)
    if value is None:
        return False
    if not isinstance(value, bool):
        raise TypeError(f'config.{GEN_IMAGE_ROW_MARKERS} must be boolean')
    return value

def resolve_generation_skip_down_sample(config, *, runtime_value=None) -> bool:
    configured = getattr(config, GENERATION_SKIP_DOWNSAMPLE, None)
    if configured is not None and (not isinstance(configured, bool)):
        raise TypeError(f'config.{GENERATION_SKIP_DOWNSAMPLE} must be boolean or null')
    runtime_configured = runtime_value is not None
    runtime_enabled = bool(runtime_value)
    if configured is None:
        return runtime_enabled
    if runtime_configured and runtime_enabled != configured:
        raise ValueError(f'SKIP_DOWN_SAMPLE conflicts with checkpoint-backed config.{GENERATION_SKIP_DOWNSAMPLE}: runtime={runtime_enabled}, checkpoint={configured}')
    return configured

def image_generation_special_ids(codebook_size: int) -> tuple[int, int, int]:
    if isinstance(codebook_size, bool) or not isinstance(codebook_size, Integral):
        raise TypeError('codebook_size must be a positive integer')
    codebook_size = int(codebook_size)
    if codebook_size < 1:
        raise ValueError('codebook_size must be a positive integer')
    return (codebook_size + ROW_MARKER_OFFSET, codebook_size + IMAGE_MASK_OFFSET, codebook_size + IMAGE_END_OFFSET)

def image_fixed_prefix_token_count(config) -> int:
    return _validate_fixed_prefix_count(getattr(config, GEN_IMAGE_FIXED_PREFIX_TOKENS, 0))

def image_generation_fixed_prefix_ids(codebook_size: int, count: int) -> tuple[int, ...]:
    image_generation_special_ids(codebook_size)
    count = _validate_fixed_prefix_count(count)
    start = int(codebook_size) + FIXED_PREFIX_OFFSET
    return tuple(range(start, start + count))

def split_generation_fixed_prefix(generation_placeholder_mask: torch.Tensor, *, codebook_size: int, count: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    if not torch.is_tensor(generation_placeholder_mask):
        raise TypeError('generation_placeholder_mask must be a tensor')
    if generation_placeholder_mask.ndim != 2:
        raise ValueError('generation_placeholder_mask must have shape [batch, sequence]')
    placeholder_mask = generation_placeholder_mask.to(dtype=torch.bool)
    fixed_ids = image_generation_fixed_prefix_ids(codebook_size, count)
    if not fixed_ids:
        return (torch.zeros_like(placeholder_mask), placeholder_mask, torch.empty(0, dtype=torch.long, device=placeholder_mask.device))
    previous_placeholder = torch.cat((torch.zeros_like(placeholder_mask[:, :1]), placeholder_mask[:, :-1]), dim=1)
    run_starts = placeholder_mask & ~previous_placeholder
    positions = torch.arange(placeholder_mask.shape[1], device=placeholder_mask.device, dtype=torch.long).unsqueeze(0).expand_as(placeholder_mask)
    most_recent_start = torch.where(run_starts, positions, positions.new_full((), -1)).cummax(dim=1).values
    run_offsets = positions - most_recent_start
    fixed_mask = placeholder_mask & run_offsets.lt(len(fixed_ids))
    image_mask = placeholder_mask & ~fixed_mask
    next_placeholder = torch.cat((placeholder_mask[:, 1:], torch.zeros_like(placeholder_mask[:, :1])), dim=1)
    run_ends = placeholder_mask & ~next_placeholder
    run_lengths = run_offsets[run_ends] + 1
    if bool(run_lengths.le(len(fixed_ids)).any()):
        shortest = int(run_lengths.min().item())
        raise ValueError(f'each generated-image placeholder run must contain the fixed prefix plus at least one image slot; shortest run={shortest}, fixed prefix={len(fixed_ids)}')
    dense_ids = int(codebook_size) + FIXED_PREFIX_OFFSET + run_offsets
    return (fixed_mask, image_mask, dense_ids[fixed_mask])

def generation_embedding_capacity(generation_model) -> int:
    embedding = getattr(generation_model, 'gen_embedding', None)
    if isinstance(embedding, nn.Embedding):
        return int(embedding.num_embeddings)
    if isinstance(embedding, nn.Module):
        for module in embedding.modules():
            if isinstance(module, nn.Embedding):
                return int(module.num_embeddings)
    raise TypeError('generation model does not expose an nn.Embedding lookup')

def image_row_marker_count(generation_model, *, skip_down_sample=False) -> int:
    compression_available = getattr(generation_model, 'packed_image_codec', None) is not None or getattr(generation_model, 'downsample_gen', None) is not None
    if compression_available and (not skip_down_sample):
        return MARKERS_PER_RAW_ROW
    return UNCOMPRESSED_MARKERS_PER_RAW_ROW

def validate_image_row_marker_support(generation_model, *, skip_down_sample=False):
    config = generation_model.config
    if 'emu3' not in 'emu3'.lower():
        raise NotImplementedError("generation image row markers currently require Emu3's single-codebook tokenizer")
    image_row_marker_count(generation_model, skip_down_sample=skip_down_sample)
    codebook_size = int(generation_model.vqvae.config.num_vq_embeddings)
    (row_id, mask_id, end_id) = image_generation_special_ids(codebook_size)
    capacity = generation_embedding_capacity(generation_model)
    if end_id >= capacity:
        raise ValueError(f'generation embedding has no spare rows for row/end markers: required ID {end_id}, capacity {capacity}')
    return (row_id, mask_id, end_id)

def validate_image_fixed_prefix_support(generation_model) -> tuple[int, ...]:
    config = generation_model.config
    count = image_fixed_prefix_token_count(config)
    if count == 0:
        return ()
    if getattr(config, 'generation_embedding_scale', None) is None:
        raise ValueError('generation fixed-prefix tokens require factorized generation scaling so uncompressed tokens receive the common input scale')
    if 'emu3' not in 'emu3'.lower():
        raise NotImplementedError("generation fixed-prefix tokens currently require Emu3's single-codebook tokenizer")
    codebook_size = int(generation_model.vqvae.config.num_vq_embeddings)
    fixed_ids = image_generation_fixed_prefix_ids(codebook_size, count)
    capacity = generation_embedding_capacity(generation_model)
    if fixed_ids[-1] >= capacity:
        raise ValueError(f'generation embedding has no spare rows for the requested fixed prefix: required ID {fixed_ids[-1]}, capacity {capacity}')
    special_ids = set(image_generation_special_ids(codebook_size))
    if special_ids.intersection(fixed_ids):
        raise RuntimeError('generation fixed-prefix IDs overlap image mask/row/end markers')
    return fixed_ids

def augmented_image_shape(gen_shape, *, markers_per_row: int=MARKERS_PER_RAW_ROW) -> tuple[int, int]:
    if gen_shape is None or len(gen_shape) != 2:
        raise ValueError('gen_shape must contain image height and width')
    markers_per_row = _validate_markers_per_row(markers_per_row)
    (height, width) = (int(value) for value in gen_shape)
    if height < 1 or width < 1:
        raise ValueError('image row markers require a positive image grid')
    if markers_per_row == MARKERS_PER_RAW_ROW and (height % 2 or width % 2):
        raise ValueError('two-column image row markers require an even-height, even-width grid for 2x spatial compression')
    return (height, width + markers_per_row)

def append_image_row_markers(token_ids: torch.Tensor, gen_shape, *, row_marker_id: int, image_end_id: int, markers_per_row: int=MARKERS_PER_RAW_ROW) -> tuple[torch.Tensor, tuple[int, int], torch.Tensor]:
    if not torch.is_tensor(token_ids) or token_ids.ndim != 2:
        raise ValueError('token_ids must have shape [batch, raw_image_tokens]')
    markers_per_row = _validate_markers_per_row(markers_per_row)
    (height, width) = (int(value) for value in gen_shape)
    augmented_shape = augmented_image_shape((height, width), markers_per_row=markers_per_row)
    if token_ids.shape[-1] != height * width:
        raise ValueError(f'raw image token count does not match gen_shape: {token_ids.shape[-1]} != {height * width}')
    grid = token_ids.reshape(token_ids.shape[0], height, width)
    markers = torch.full((token_ids.shape[0], height, markers_per_row), int(row_marker_id), dtype=token_ids.dtype, device=token_ids.device)
    markers[:, -1, :] = int(image_end_id)
    augmented = torch.cat((grid, markers), dim=-1).flatten(1)
    structural_mask = torch.zeros_like(augmented, dtype=torch.bool)
    structural_mask.reshape(token_ids.shape[0], *augmented_shape)[:, :, -markers_per_row:] = True
    return (augmented, augmented_shape, structural_mask)

def strip_image_row_markers(token_ids: torch.Tensor, original_gen_shape, *, markers_per_row: int=MARKERS_PER_RAW_ROW) -> torch.Tensor:
    if not torch.is_tensor(token_ids) or token_ids.ndim != 2:
        raise ValueError('token_ids must have shape [batch, augmented_tokens]')
    markers_per_row = _validate_markers_per_row(markers_per_row)
    (height, width) = (int(value) for value in original_gen_shape)
    augmented_shape = augmented_image_shape((height, width), markers_per_row=markers_per_row)
    if token_ids.shape[-1] != math.prod(augmented_shape):
        raise ValueError(f'augmented image token count does not match marker grid: {token_ids.shape[-1]} != {math.prod(augmented_shape)}')
    return token_ids.reshape(token_ids.shape[0], *augmented_shape)[:, :, :width].flatten(1)