"""Strict phase-policy checkpoint loading, retaining the frozen older factories."""
import io
from pathlib import Path

import torch

from tinyasr.codec_factory_v2 import load_codec_checkpoint as load_previous
from tinyasr.codec_phase_residual_v1 import ARCHITECTURE, POLICIES, PhaseResidualCodec


def load_codec_checkpoint(path):
    # A fresh stream is needed by older factories that reopen their input.
    if isinstance(path, (str, Path)):
        with Path(path).open('rb') as source:
            raw = source.read(100_000_001)
    else:
        raw = path.read(100_000_001)
    if not isinstance(raw, bytes) or not 0 < len(raw) <= 100_000_000:
        raise ValueError('Invalid checkpoint bytes')
    state = torch.load(io.BytesIO(raw), map_location='cpu', weights_only=False)
    if state.get('architecture') != ARCHITECTURE:
        # Factory v2 reopens fallback checkpoints; v1 must receive a fresh stream.
        from tinyasr.codec_projected_warmup_v1 import ARCHITECTURE as PROJECTED
        from tinyasr.codec_factory_v1 import load_codec_checkpoint as load_v1
        return (load_previous if state.get('architecture') == PROJECTED else load_v1)(io.BytesIO(raw))
    policy = state.get('phase_policy')
    if policy not in POLICIES or state.get('config', {}).get('phase_policy', policy) != policy:
        raise ValueError('Missing or inconsistent phase policy')
    model = PhaseResidualCodec(policy)
    model.load_state_dict(state['codec'], strict=True)
    return model.eval().requires_grad_(False)