"""102 payload bits -> turbo -> CTR mask -> ten binary-latent glyphs.

Analysis always crops the saved, quantized BMP; latent vectors and channel bits
are never retained. The reference secret is used only AFTER turbo decoding.
"""
from collections import OrderedDict
import base64
import io
import math
import secrets
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from analyzer import TransposedStegoAnalyzer
from ccsds_turbo_102_640 import encode, decode_llr
from renderer import LetterRenderer, validate_text, tile_layout, MARGIN

AMPLITUDE = math.sqrt(2 / math.pi)


def secret_bits(value):
    if not isinstance(value, str) or not value.isascii():
        raise ValueError('Secret must contain ASCII characters only (each stored as 8 bits).')
    return np.unpackbits(np.frombuffer(value.encode('ascii'), dtype=np.uint8), bitorder='big')


def bits_secret(bits):
    bits = np.asarray(bits, dtype=np.uint8).ravel()
    # Ignore only the trailing partial byte used to complete the turbo blocks.
    # Latin-1 preserves all recovered byte values, including erroneous high bits.
    return np.packbits(bits[:len(bits) // 8 * 8], bitorder='big').tobytes().decode('latin-1')


def payload_bits(value, count):
    return np.pad(secret_bits(value), (0, count - len(value) * 8))


def ctr_bits(key, iv, count):
    for name, value in [('Key', key), ('IV', iv)]:
        if not isinstance(value, str) or len(value) != 16 or not value.isascii():
            raise ValueError(f'{name} must be exactly 16 ASCII characters (16 bytes).')
    if count % 8:
        raise ValueError('CTR bit count must be byte aligned.')
    cipher = Cipher(algorithms.AES(key.encode('ascii')), modes.CTR(iv.encode('ascii'))).encryptor()
    stream = cipher.update(bytes(count // 8)) + cipher.finalize()
    return np.unpackbits(np.frombuffer(stream, dtype=np.uint8), bitorder='big')


def prepare(payload):
    if not isinstance(payload, dict):
        raise ValueError('Request must be a JSON object.')
    text = validate_text(payload.get('text'))
    count = sum(c not in ' \n' for c in text)
    suffix = 'ENDOFLETTER'[:(-count) % 10]
    text = validate_text(text + suffix)
    count += len(suffix)
    capacity = (count // 10 * 102) // 8
    secret = payload.get('secret', '')
    secret_bits(secret)
    if len(secret) > capacity:
        raise ValueError(f'Secret exceeds capacity: {len(secret)} ASCII characters supplied, {capacity} fit.')
    tile = payload.get('tile_size', 28)
    if type(tile) is not int or tile not in (28, 56):
        raise ValueError('Use 28 px or 56 px tiles to preserve native glyph pixels.')
    return text, suffix, count, capacity, secret, secret.ljust(capacity, '\0')


class StegoApp:
    def __init__(self, root: Path, device='auto'):
        self.root = root
        self.renderer = LetterRenderer(root / 'models/generator.pt', device)
        self.sessions = OrderedDict()
        self.models = OrderedDict()

    def info(self):
        return {**self.renderer.info(), 'analyzers': 36, 'payload_bits': 102, 'channel_bits': 640}

    def model(self, label):
        if label not in self.models:
            checkpoint = torch.load(self.root / 'models/analyzers' / (label + '.pt'), map_location='cpu', weights_only=True)
            model = TransposedStegoAnalyzer()
            model.load_state_dict(checkpoint['model'])
            self.models[label] = model.to(self.renderer.device).eval()
            # Bounded GPU/CPU model cache, independent of the number of labels.
            if len(self.models) > 4:
                self.models.popitem(last=False)
        self.models.move_to_end(label)
        return self.models[label]

    def generate(self, payload):
        text, suffix, count, capacity, original, padded = prepare(payload)
        mask = ctr_bits(payload.get('key'), payload.get('iv'), count * 64)
        data = payload_bits(padded, count // 10 * 102).reshape(-1, 102)
        coded = np.concatenate([encode(block) for block in data])
        channel = coded ^ mask
        latents = torch.from_numpy((channel.astype(np.float32) * 2 - 1).reshape(-1, 64) * AMPLITUDE)
        bitmap, details = self.renderer.render({**payload, 'text': text, '_latents': latents})
        tile = payload.get('tile_size', 28)
        gap = payload.get('line_spacing', 12)
        rows = tile_layout(text, details['columns'])
        glyphs = [(char, MARGIN + col * tile, MARGIN + row * (tile + gap))
                  for row, line in enumerate(rows) for col, char in enumerate(line) if char != ' ']
        token = secrets.token_urlsafe(24)
        self.sessions[token] = {'bitmap': bitmap, 'glyphs': glyphs, 'tile': tile,
                                'reference': padded, 'original_length': len(original)}
        while len(self.sessions) > 2:
            self.sessions.popitem(last=False)
        return {'letter_id': token, 'bmp': base64.b64encode(bitmap).decode('ascii'),
                'text': text, 'suffix': suffix, 'capacity': capacity, 'secret_padding': capacity - len(original),
                **{k:v for k,v in details.items() if k != 'seed'}}

    @torch.inference_mode()
    def analyze(self, payload):
        if not isinstance(payload, dict):
            raise ValueError('Request must be a JSON object.')
        token = payload.get('letter_id')
        if not isinstance(token, str) or token not in self.sessions:
            raise ValueError('Generate a letter first. Old letters expire after two generations or server restart.')
        session = self.sessions[token]
        glyphs = session['glyphs']
        mask = ctr_bits(payload.get('key'), payload.get('iv'), len(glyphs) * 64)
        page = Image.open(io.BytesIO(session['bitmap'])).convert('L')
        logits = np.empty((len(glyphs), 64), dtype=np.float32)
        labels = sorted(set(g[0] for g in glyphs))
        tile = session['tile']
        for label in labels:
            indices = [i for i, g in enumerate(glyphs) if g[0] == label]
            model = self.model(label)
            for start in range(0, len(indices), 64):
                selected = indices[start:start+64]
                images = []
                for i in selected:
                    _, x, y = glyphs[i]
                    crop = np.asarray(page.crop((x, y, x+tile, y+tile)))
                    # Integer nearest-neighbor enlargement is exactly reversible.
                    crop = crop[::tile//28, ::tile//28]
                    images.append(1 - crop.astype(np.float32) / 127.5)
                batch = torch.from_numpy(np.stack(images)[:, None]).to(self.renderer.device)
                logits[selected] = model(batch).cpu().numpy()
        # Network logit = log P(1)/P(0); codec LLR = log P(0)/P(1).
        llr = -logits.ravel().astype(np.float64) * (1 - 2 * mask.astype(np.float64))
        recovered = np.concatenate([decode_llr(block) for block in llr.reshape(-1, 640)])
        decoded = bits_secret(recovered)
        # Reference becomes available only after inference and decoding finish.
        reference = session['reference']
        expected = payload_bits(reference, len(recovered))
        bit_errors = int(np.count_nonzero(expected != recovered))
        char_errors = sum(a != b for a, b in zip(decoded, reference))
        block_errors = [int(np.count_nonzero(a != b)) for a,b in zip(expected.reshape(-1,102), recovered.reshape(-1,102))]
        original_length = session['original_length']
        return {'secret': decoded, 'original_secret': decoded[:original_length],
                'padding_symbols': len(decoded)-original_length, 'bit_errors': bit_errors,
                'character_errors': char_errors, 'total_bits': len(recovered),
                'total_symbols': len(decoded), 'block_bit_errors': block_errors,
                'correct': bit_errors == 0}