"""Uploads for hosted sessions; no access to server filesystem paths."""
from pathlib import Path

import streamlit as st

from app_environment import UPLOAD_ROOT, new_session_directory, remove_session_directory


MAX_FILE_BYTES = 1024 * 1024 * 1024
MAX_TOTAL_BYTES = 2 * MAX_FILE_BYTES


def store_uploads(files, tokenizer=None, root=UPLOAD_ROOT):
    if not files:
        raise ValueError('Choose a GGUF or safetensors file.')
    names = [item.name for item in files]
    if any(not name or '/' in name or '\\' in name or ':' in name or '\x00' in name or Path(name).name != name for name in names):
        raise ValueError('Use plain file names for uploaded model files.')
    if len(set(name.casefold() for name in names)) != len(names):
        raise ValueError('Uploaded model files must have different names.')
    if any(Path(name).suffix.lower() not in ('.gguf', '.safetensors') for name in names):
        raise ValueError('Upload GGUF or safetensors weights only.')
    if any(Path(name).suffix.lower() == '.gguf' for name in names) and len(files) != 1:
        raise ValueError('Upload one GGUF file, or one or more safetensors shards.')
    all_files = list(files) + ([tokenizer] if tokenizer is not None else [])
    sizes = [len(item.getbuffer()) for item in all_files]
    if any(size > MAX_FILE_BYTES for size in sizes) or sum(sizes) > MAX_TOTAL_BYTES:
        raise ValueError('Uploads allow up to 1 GB per file and 2 GB in total. Use Hugging Face for larger models.')
    directory = new_session_directory(root)
    try:
        for item in files:
            with (directory / item.name).open('wb') as target:
                target.write(item.getbuffer())
        tokenizer_path = None
        if tokenizer is not None:
            tokenizer_path = directory / 'tokenizer.json'
            with tokenizer_path.open('wb') as target:
                target.write(tokenizer.getbuffer())
        model_path = directory / names[0] if len(names) == 1 else directory
        return model_path, tokenizer_path, directory
    except Exception:
        remove_session_directory(directory, root)
        raise


def render_upload_loader(activate_model):
    st.caption('Upload a single GGUF file or safetensors weights. Up to 1 GB per file, 2 GB total.')
    files = st.file_uploader('Model files', type=['gguf', 'safetensors'], accept_multiple_files=True, key='uploaded_models')
    tokenizer = st.file_uploader('Matching tokenizer.json (optional)', type=['json'], key='uploaded_tokenizer')
    st.caption('Uploads are processed by this Space in your session and removed when you clear the model. Restarting the Space also clears uploads.')
    if st.button('Load uploaded model', key='load_upload', type='primary', disabled=not files, use_container_width=True):
        directory = None
        try:
            with st.spinner('Reading uploaded model…'):
                path, tokenizer_path, directory = store_uploads(files, tokenizer)
                activate_model(path, tokenizer_path, 'Uploaded: ' + ', '.join(item.name for item in files))
            st.session_state.upload_directory = directory
            st.rerun()
        except Exception as error:
            if directory is not None:
                remove_session_directory(directory, UPLOAD_ROOT)
            st.error(f'Could not load upload: {error}')