"""Strict SkinTokens worker: prewarmed model, separate Blender process, preserved input."""
from pathlib import Path
import hashlib
import json
import shutil
import uuid
import time
import traceback

import spaces
import gradio as gr
import torch

from prepare import prepare, LOCK, ROOT
from preservation import verify

prepare()
from src.data.dataset import DatasetConfig, RigDatasetModule  # noqa: E402
from src.data.transform import Transform  # noqa: E402
from src.tokenizer.parse import get_tokenizer  # noqa: E402
from src.server.spec import BpyClient, bpy_server_request, get_model  # noqa: E402

# Blender CPU imports and all CUDA model placement happen before a user books a GPU.
BPY = BpyClient.launch()
MODEL = get_model(LOCK['checkpoint']).eval()
TOKENIZER = get_tokenizer(**MODEL.tokenizer_config)
TRANSFORM = Transform.parse(**MODEL.transform_config['predict_transform'])


def sha256(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


@spaces.GPU(duration=240)
def predict(source):
    config = DatasetConfig.parse(
        shuffle=False, batch_size=1, num_workers=0, pin_memory=False, persistent_workers=False,
        datapath={'data_name': None, 'loader': 'bpy_server', 'filepaths': {'articulation': [str(source)]}},
    ).split_by_cls()
    module = RigDatasetModule(predict_dataset_config=config, predict_transform=TRANSFORM,
                              tokenizer=TOKENIZER, process_fn=MODEL._process_fn)
    batch = next(iter(module.predict_dataloader()['articulation']))
    batch = {key: value.to(MODEL.device) if isinstance(value, torch.Tensor) else value
             for key, value in batch.items()}
    batch.pop('skeleton_tokens', None)
    batch.pop('skeleton_mask', None)
    batch['generate_kwargs'] = dict(max_length=2048, top_k=5, top_p=.95, temperature=1.0,
        repetition_penalty=2.0, num_return_sequences=1, num_beams=10, do_sample=True)
    with torch.inference_mode():
        asset = MODEL.predict_step(batch, skeleton_tokens=None, make_asset=True)['results'][0].asset
    if asset is None or asset.skin is None:
        raise RuntimeError('SkinTokens produced no skinned asset')
    # Asset stores NumPy data; no CUDA tensors leave this worker.
    return asset


def rig(source):
    global BPY
    output = ROOT / 'outputs' / ('rig-' + uuid.uuid4().hex)
    output.mkdir(parents=True)
    original = output / 'source.glb'
    shutil.copyfile(source, original)
    receipt = {'provider': 'native SkinTokens', 'lock': LOCK, 'source_sha256': sha256(original),
               'started': time.time(), 'stage': 'load', 'status': 'running'}
    try:
        # Surface a Blender/input failure before requesting the GPU.
        if not BPY.is_alive():
            BPY = BpyClient.launch()
        bpy_server_request('load', str(original))
        receipt['stage'] = 'predict'
        asset = predict(original)
        receipt['stage'] = 'transfer'
        result = output / 'rigged.glb'
        transfer = bpy_server_request('transfer', dict(source_asset=asset, target_path=str(original),
                                           export_path=str(result), group_per_vertex=4))
        if transfer != 'ok':
            raise RuntimeError(f'Blender transfer failed: {transfer}')
        receipt['stage'] = 'verify'
        receipt['preservation'] = verify(original, result)
        receipt.update(status='complete', output_sha256=sha256(result), finished=time.time())
        return str(result), receipt
    except Exception as error:
        traceback.print_exc()
        message = f"{receipt['stage']}: {type(error).__name__}: {error}"
        receipt.update(status='failed', error=message, finished=time.time())
        raise gr.Error(message) from error
    finally:
        (output / 'provenance.json').write_text(json.dumps(receipt, indent=2))


with gr.Blocks(title='Component Studio SkinTokens') as demo:
    source = gr.File(file_types=['.glb'], type='filepath', label='Textured GLB')
    button = gr.Button('Rig character')
    result = gr.File(label='Rigged GLB')
    report = gr.JSON(label='Preservation and provenance')
    button.click(rig, source, [result, report], api_name='rig', concurrency_limit=1)

demo.queue(max_size=4).launch(server_name='0.0.0.0', server_port=7860,
                            allowed_paths=[str(ROOT / 'outputs')])