import io
import pickle
import torch
import streamlit as st
import config


class _CPUUnpickler(pickle.Unpickler):
    def find_class(self, module, name):
        if module == 'torch.storage' and name == '_load_from_bytes':
            return lambda b: torch.load(io.BytesIO(b), map_location='cpu', weights_only=False)
        return super().find_class(module, name)


def _bundle_source(filename: str):
    local_path = config.MODELS_DIR / filename
    if local_path.is_file():
        return open(local_path, 'rb')
    model_bytes = st.session_state.get('model_bytes', {})
    if filename in model_bytes:
        return io.BytesIO(model_bytes[filename])
    return None


def _is_available(filename: str) -> bool:
    if (config.MODELS_DIR / filename).is_file():
        return True
    return filename in st.session_state.get('model_bytes', {})


def _load_bundle(filename: str) -> dict | None:
    cache_key = f'_bundle_cache_{filename}'
    if cache_key in st.session_state:
        return st.session_state[cache_key]
    src = _bundle_source(filename)
    if src is None:
        return None
    with src as f:
        st.session_state[cache_key] = _CPUUnpickler(f).load()
    return st.session_state[cache_key]


def load_bundle(arch: str, data_type: str, target: str) -> dict | None:
    safe = data_type.replace('+', '_plus_')
    return _load_bundle(f'final_{arch}_{safe}_{target}.pkl')

def load_ensemble_bundle(variant: str, target: str) -> dict | None:
    return _load_bundle(f'final_ensemble_{variant}_{target}.pkl')


def model_status(arch: str) -> dict[str, bool]:
    result = {}
    for dt in config.DATA_TYPES:
        safe = dt.replace('+', '_plus_')
        result[dt] = all(
            _is_available(f'final_{arch}_{safe}_{t}.pkl')
            for t in config.TARGETS
        )
    return result