"""Optional CUDA compilation; exact trained weights and decision contract.

Compilation is an explicit runtime choice. Cold compilation is substantial;
hosts should warm the five supported sequence buckets before accepting audio.
"""
from .router import Router


BUCKETS = (32, 64, 128, 256, 384)


def _cudagraph_failure(error):
    """Identify graph-runtime failures without inspecting request data."""
    seen = set()
    while error is not None and id(error) not in seen:
        seen.add(id(error))
        trace = error.__traceback__
        while trace is not None:
            module = trace.tb_frame.f_globals.get('__name__', '')
            if module.startswith('torch._inductor.cudagraph'):
                return True
            trace = trace.tb_next
        error = error.__cause__ or error.__context__
    return False


class _BucketedModel:
    """Keep graph failure recovery under the caller's existing device admission."""

    def __init__(self, native, compiled, pad_id):
        import torch

        self.native = native
        self.compiled = compiled
        self.pad_id = pad_id
        self.thread_budget = torch.get_num_threads()
        self.execution_mode = 'compiled-cuda'

    def _recover(self, input_ids, attention_mask, observed, reason):
        # Validate native execution before committing the persistent fallback.
        output = self.native(input_ids, attention_mask, observed)
        self.execution_mode = 'eager-cuda-fallback'
        print(f'[jev] compiled CUDA unavailable ({reason}); '
              'continuing with same-weight eager CUDA inference', flush=True)
        return output

    def __call__(self, input_ids, attention_mask, observed):
        import torch
        from torch.nn import functional as F

        if input_ids.shape[0] != 2 or observed.shape != (1, 7):
            raise ValueError('Compiled router is a single-decision runtime')
        length = input_ids.shape[1]
        bucket = next((n for n in BUCKETS if n >= length), None)
        if bucket is None:
            raise ValueError('Sequence exceeds decision contract')
        if bucket > length:
            input_ids = F.pad(input_ids, (0, bucket-length), value=self.pad_id)
            attention_mask = F.pad(attention_mask, (0, bucket-length), value=0)
        # Match the layout used by prewarming, including a size-one dimension.
        observed = observed.reshape(7).unsqueeze(0)
        if self.execution_mode == 'compiled-cuda':
            # A provider can change this setting on an input worker. Avoid an
            # expensive recompilation before the known graph TLS failure.
            if torch.get_num_threads() != self.thread_budget:
                return self._recover(input_ids, attention_mask, observed, 'thread-budget-changed')
            try:
                torch.compiler.cudagraph_mark_step_begin()
                return self.compiled(input_ids, attention_mask, observed)
            except Exception as error:
                if isinstance(error, torch.OutOfMemoryError) or not _cudagraph_failure(error):
                    raise
                return self._recover(input_ids, attention_mask, observed, type(error).__name__)
        return self.native(input_ids, attention_mask, observed)


class CompiledCudaRouter(Router):
    def __init__(self, path, *, device='cuda:0', admission=None):
        if not device.startswith('cuda'):
            raise ValueError('Compiled runtime is CUDA-only')
        super().__init__(path, backend='eager', device=device, admission=admission)

    def _load(self):
        import torch
        super()._load()
        native = self.model
        compiled = torch.compile(native, mode='reduce-overhead', dynamic=False)
        pad_id = native.encoder.config.pad_token_id
        self.model = _BucketedModel(native, compiled, pad_id)