import spaces
from spaces.config import Config as SpacesConfig

from collections import deque
import base64
import os
from pathlib import Path
import sys
import tempfile
import threading
import time
import uuid

import gradio as gr
from huggingface_hub import hf_hub_download
import numpy as np
import torch

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT / 'space'))
from space.assets_loader import resolve_t5, REVISION, REPO
from space.inference import load_model
from space.text_encoder import load_text_encoder, load_prompt_features, save_prompt_features
from fast_window_runtime import FastWindowRuntime
from hml263_runtime import load_hml263_model
from window_control import WindowPath as SeedPath
from hml263_control import WindowPath as HumanML3DPath
from window_viewer import create_viewer
from cloud_sessions import SessionStore

MODELS = ('SEED', 'HumanML3D')
DEFAULTS = [('A person walks forward.', .8),
            ('A person runs forward.', 2.5),
            ('A person walks forward in a crouched position.', .45),
            ('A person is dancing.', .8)]
GPU_CONCURRENCY = None if SpacesConfig.zero_gpu else 1
LOCAL_GPU_LOCK = threading.Lock()
STORE = SessionStore(Path(tempfile.gettempdir()) / 'flood2_live_window_controls')
FEATURES = ROOT / 'space/assets/prompts.npz'
DEFAULT_BANK = load_prompt_features(ROOT / 'space/assets/default_prompts.npz')
BASE_BANK = load_prompt_features(FEATURES)
seed_checkpoint = os.getenv('FLOOD2_SEED_CHECKPOINT') or hf_hub_download(
    REPO, 'checkpoints/seed_path_fk_300k/model.ckpt', revision=REVISION)
hml_checkpoint = os.getenv('FLOOD2_HML_CHECKPOINT') or hf_hub_download(
    REPO, 'checkpoints/humanml3d_babel_path_200k/model.ckpt', revision=REVISION)
hml_config = os.getenv('FLOOD2_HML_CONFIG') or hf_hub_download(
    REPO, 'checkpoints/humanml3d_babel_path_200k/config.yaml', revision=REVISION)
RUNTIMES = {
    'SEED': FastWindowRuntime.from_runtime(load_model(ROOT/'space/vendor', seed_checkpoint, FEATURES)),
    'HumanML3D': load_hml263_model(ROOT/'space/vendor', hml_checkpoint, hml_config, FEATURES),
}
# Module-scope loads are virtualized by ZeroGPU. Encoding and graph capture
# happen only inside decorated allocations.
encoder_path, tokenizer_path = resolve_t5()
ENCODER = load_text_encoder(ROOT/'space/vendor', encoder_path, tokenizer_path)


def owner(request):
    value = getattr(request, 'session_hash', None)
    if not value:
        raise gr.Error('Reload the page to start a session.')
    return value


def checked_model(value):
    if value not in MODELS:
        raise gr.Error('Choose SEED or HumanML3D.')
    return value


def read(session, request):
    try:
        return STORE.read(session, owner(request))
    except (ValueError, FileNotFoundError) as exc:
        raise gr.Error(str(exc)) from exc


@spaces.GPU(duration=30)
def prepare(previous_session, model, p1, p2, p3, p4, v1, v2, v3, v4, request: gr.Request):
    model = checked_model(model)
    browser = owner(request)
    prompts = [str(prompt).strip() for prompt in (p1, p2, p3, p4)]
    if any(not prompt or len(prompt) > 500 for prompt in prompts):
        raise gr.Error('Enter 1 to 500 characters in every prompt slot.')
    try:
        speeds = [float(value) for value in (v1, v2, v3, v4)]
    except (TypeError, ValueError):
        raise gr.Error('Every speed must be between 0 and 5 m/s.')
    if any(not np.isfinite(value) or not 0 <= value <= 5 for value in speeds):
        raise gr.Error('Every speed must be between 0 and 5 m/s.')
    if previous_session:
        previous = read(previous_session, request)
        if previous['state'] not in ('paused', 'error') or previous.get('stream_active'):
            raise gr.Error('Stop before starting a new motion.')
    bank = {text: DEFAULT_BANK[text] for text in set(prompts) if text in DEFAULT_BANK}
    missing = [text for text in dict.fromkeys(prompts) if text not in bank]
    if missing:
        bank.update(ENCODER.encode(missing))
    session = uuid.uuid4().hex
    save_prompt_features(STORE.path(session).with_suffix('.npz'), bank)
    STORE.create(session, dict(owner=browser, model=model, slots=prompts, speeds=speeds,
                              slot=0, x=0., z=0., client_seq=-1, keys_updated_at=0.,
                              version=0, state='starting', stop=False, stream_active=False,
                              frames=0, runtime={}, error=None))
    return session


def control(session, seq, x, z, shift=False, slot=0, request: gr.Request=None):
    if not session:
        return 'Ready'
    try:
        data = STORE.control(session, owner(request), seq, x, z, slot)
    except (ValueError, FileNotFoundError) as exc:
        raise gr.Error(str(exc)) from exc
    return 'Running' if data['state'] == 'running' else 'Ready'


def keyboard_control(event: gr.EventData, request: gr.Request):
    data = event._data
    if isinstance(data, dict):
        control(data.get('session_id', ''), data.get('seq', -1), data.get('x', 0.),
                data.get('z', 0.), False, data.get('slot', 0), request)


def stop_session(session, request: gr.Request):
    if not session:
        return 'Ready'
    try:
        data = STORE.stop(session, owner(request))
    except (ValueError, FileNotFoundError) as exc:
        raise gr.Error(str(exc)) from exc
    return 'Stopping…' if data['stream_active'] else 'Stopped'


def session_status(session, request: gr.Request):
    if not session:
        return dict(state='ready', frames=0, stream_active=False)
    data = read(session, request)
    return {key: value for key, value in data.items() if key != 'owner'}


@spaces.GPU(duration=60)
def stream(session, request: gr.Request):
    browser = owner(request)
    initial = STORE.claim(session, browser)
    model = initial['model']
    runtime = RUNTIMES[model]
    planner = SeedPath() if model == 'SEED' else HumanML3DPath()
    fps = 30 if model == 'SEED' else 20
    completed = sequence = 0
    started = time.perf_counter()
    deadline = started + 55.
    playback_origin = first_frame_seconds = None
    compute_seconds = 0.
    frame_seconds = deque(maxlen=60)
    text_metas = []
    latest = initial
    last_meta = dict(slot=0, frame=-1, prompt=initial['slots'][0])
    gpu_name = None
    acquired = False
    runtime_started = False

    def packet(status, frames=None, metas=None, rotations=None, error=None):
        nonlocal sequence
        frames, metas = frames or [], metas or []
        info = runtime.status() if runtime_started else {}
        active_slot = last_meta['slot'] if completed else latest['slot']
        result = dict(session_id=session, sequence=sequence, run_revision=1,
                      model=model, representation='mesh' if model == 'SEED' else 'skeleton',
                      start_frame=completed-len(frames), frames=frames, frame_meta=metas,
                      control_points=[dict(meta) for meta in metas]+[dict(point) for point in planner.points],
                      control_points_delta=True, control_replace_from_frame=completed-len(frames),
                      spring_point=dict(planner.points[-1]) if planner.points else None,
                      target_point=planner.target_point(active_slot),
                      slots=list(latest['slots']), speeds=list(latest['speeds']),
                      active_slot=active_slot, base_slot=latest['slot'],
                      prompt=latest['slots'][active_slot], action=latest['slots'][active_slot],
                      speed=latest['speeds'][active_slot], completed_frame=completed-1,
                      input_frame=info.get('steps', 0)-1, fps=fps, generated_frames=completed,
                      elapsed_seconds=time.perf_counter()-started, first_frame_seconds=first_frame_seconds,
                      inference_fps=len(frame_seconds)/max(sum(frame_seconds), 1e-9),
                      session_inference_fps=completed/max(compute_seconds, 1e-9),
                      graph_enabled=bool(info.get('graph', {}).get('replays')),
                      graph=info.get('graph'), text_bank=info.get('fixed_text'),
                      control_version=last_meta.get('version', 0), path_mode='window',
                      prompt_history=dict(enabled=True, max_history=30,
                                          resets=info.get('history_resets', 0),
                                          visible_history=info.get('visible_history', 0),
                                          cutoff=info.get('history_cutoff_absolute')),
                      gpu=gpu_name, status=status, error=error)
        if rotations:
            result['mesh_pose'] = dict(
                rotations=base64.b64encode(np.asarray(rotations, dtype='<f4').tobytes()).decode('ascii'),
                origins=[frame[0] for frame in frames], frame_count=len(rotations))
        STORE.patch(session, browser, frames=completed, runtime=info, error=error)
        sequence += 1
        return result

    try:
        if not SpacesConfig.zero_gpu:
            LOCAL_GPU_LOCK.acquire()
            acquired = True
        gpu_name = torch.cuda.get_device_name()
        # A reused worker always starts with a fresh latent, planner, decoder,
        # text bank and captured graph; no mutable state crosses sessions.
        runtime.close()
        runtime.model.text_module.text_cache = {}
        runtime.add_prompt_features(BASE_BANK)
        runtime.add_prompt_features(load_prompt_features(STORE.path(session).with_suffix('.npz')))
        yield packet('start')
        latest = STORE.read(session, browser)
        if latest['stop'] or time.perf_counter() >= deadline:
            STORE.patch(session, browser, state='paused', stop=True, stream_active=False)
            yield packet('paused')
            return
        tick = time.perf_counter()
        runtime.start(seed=0, history=60, use_graph=True,
                      reset_history_on_prompt_change=True, cfg_scale=4.)
        runtime_started = True
        compute_seconds += time.perf_counter()-tick
        latest = STORE.read(session, browser)
        if not latest['stop']:
            latest = STORE.patch(session, browser, state='running')
        while time.perf_counter() < deadline:
            latest = STORE.read(session, browser)
            if latest['stop']:
                break
            tick = time.perf_counter()
            alive = time.time()-latest['keys_updated_at'] < 1.2
            x, z = (latest['x'], latest['z']) if alive else (0., 0.)
            slot = latest['slot']
            prompt = latest['slots'][slot]
            for meta in text_metas:
                meta.update(slot=slot, base_slot=slot, run=False,
                            prompt=prompt, text_version=latest['version'])
            text_metas.append(dict(frame=runtime.status()['steps'], slot=slot, base_slot=slot,
                                   run=False, prompt=prompt, text_version=latest['version']))
            rows, points = planner.plan(text_metas, x=x, z=z, shift=False, slot=slot,
                                        speeds=latest['speeds'], version=latest['version'])
            raw = runtime.step_replanned(prompt, rows, replan_text=True)
            frames, metas, rotations = [], [], []
            if raw is not None:
                native, meta = planner.commit_first()
                text_metas.pop(0)
                if meta['frame'] != completed or not np.allclose(raw[:3], native, rtol=1e-5, atol=1e-6):
                    raise RuntimeError('Motion and path frame alignment was lost.')
                last_meta = meta
                frames = [np.round(runtime.recover(raw)[0], 5).tolist()]
                metas = [dict(meta)]
                if model == 'SEED':
                    rotations = [runtime.render_rotations(raw)]
                completed += 1
                if first_frame_seconds is None:
                    first_frame_seconds = time.perf_counter()-started
                    playback_origin = time.perf_counter()
            seconds = time.perf_counter()-tick
            compute_seconds += seconds
            if raw is not None:
                frame_seconds.append(seconds)
            if frames or runtime.status()['steps'] % 6 == 0:
                yield packet('running' if completed else 'start', frames, metas, rotations)
            if playback_origin is not None:
                delay = playback_origin+completed/fps-time.perf_counter()
                if delay > 0:
                    time.sleep(min(delay, .04))
        STORE.patch(session, browser, state='paused', stop=True, stream_active=False,
                    x=0., z=0., keys_updated_at=0.)
        yield packet('paused')
    except GeneratorExit:
        raise
    except Exception as exc:
        import traceback
        traceback.print_exc()
        message = f'{type(exc).__name__}: {exc}'
        STORE.patch(session, browser, state='error', stop=True, stream_active=False, error=message)
        yield packet('error', error=message)
    finally:
        try:
            runtime.close()
        finally:
            try:
                state = STORE.read(session, browser)
                STORE.patch(session, browser, state='error' if state['error'] else 'paused',
                            stop=True, stream_active=False, x=0., z=0., keys_updated_at=0.)
            finally:
                if acquired:
                    LOCAL_GPU_LOCK.release()


with gr.Blocks(title='FloodDiffusion 2') as demo:
    gr.Markdown('# FloodDiffusion 2')
    session = gr.Textbox(value='', visible=False)
    body_loaded = gr.State(False)
    preparing = gr.State(False)
    model = gr.Radio(MODELS, value='SEED', show_label=False, container=False)
    with gr.Row():
        start = gr.Button('Start', variant='primary', interactive=False)
        stop = gr.Button('Stop', interactive=False)
    prompts, speeds = [], []
    with gr.Accordion('Prompts & speeds', open=False):
        for index, (text, value) in enumerate(DEFAULTS, 1):
            with gr.Row():
                prompts.append(gr.Textbox(value=text, label=str(index), lines=1, max_length=500, scale=5))
                speeds.append(gr.Number(value=value, label='Speed (m/s)', minimum=0,
                                        maximum=5, step=.05, scale=1, min_width=120))
    status = gr.Markdown('Loading…')
    viewer = create_viewer(require_body_ready=True, action_slots=True)

    def sync_controls(session_id, loaded, preparing_now, request: gr.Request):
        info = session_status(session_id, request)
        phase = info['state']
        busy = bool(preparing_now) or phase in ('starting', 'running', 'pausing')
        message = {'ready':'', 'paused':'Stopped', 'starting':'Preparing…',
                   'running':'', 'pausing':'Stopping…', 'error':info.get('error', '')}.get(phase, '')
        if not loaded:
            message = 'Loading…'
        elif preparing_now:
            message = 'Preparing…'
        return (gr.Button(interactive=bool(loaded) and not busy and phase in ('ready', 'paused', 'error')),
                gr.Button(interactive=busy and not preparing_now and phase != 'pausing'),
                gr.Radio(interactive=not busy),
                *(gr.Textbox(interactive=not busy) for _ in prompts),
                *(gr.Number(interactive=not busy) for _ in speeds), message)

    ui_outputs = [start, stop, model, *prompts, *speeds, status]
    loaded = viewer.body_ready(lambda: True, outputs=body_loaded, queue=False, api_name=False)
    ui_inputs = [session, body_loaded, preparing]
    loaded.then(sync_controls, ui_inputs, ui_outputs, queue=False, api_name=False)
    gr.Timer(.5).tick(sync_controls, ui_inputs, ui_outputs,
                       queue=False, trigger_mode='always_last', api_name=False)
    def lock_controls():
        return (gr.Button(interactive=False), gr.Button(interactive=False),
                gr.Radio(interactive=False),
                *(gr.Textbox(interactive=False) for _ in prompts),
                *(gr.Number(interactive=False) for _ in speeds), 'Preparing…', True)
    locking = start.click(lock_controls, outputs=[*ui_outputs, preparing], queue=False, api_name=False)
    beginning = locking.success(prepare, [session, model, *prompts, *speeds], session,
                                api_name='prepare', concurrency_limit=GPU_CONCURRENCY,
                                concurrency_id='gpu' if GPU_CONCURRENCY else None)
    prepared = beginning.success(lambda: False, outputs=preparing, queue=False, api_name=False)
    prepared.success(sync_controls, ui_inputs, ui_outputs, queue=False, api_name=False)
    running = prepared.success(stream, session, viewer, api_name='stream',
                                 concurrency_limit=GPU_CONCURRENCY, stream_every=.05,
                                 concurrency_id='gpu' if GPU_CONCURRENCY else None)
    running.then(sync_controls, ui_inputs, ui_outputs, queue=False, api_name=False)
    failed = beginning.failure(lambda: False, outputs=preparing, queue=False, api_name=False)
    failed.then(sync_controls, ui_inputs, ui_outputs, queue=False, api_name=False)
    stop.click(stop_session, session, status, queue=False, api_name='stop')
    viewer.control(keyboard_control, inputs=[], outputs=[], queue=False, trigger_mode='multiple', api_name=False)
    with gr.Row(visible=False):
        control_sequence = gr.Number(value=0)
        control_x, control_z = gr.Number(value=0), gr.Number(value=0)
        control_shift, control_slot = gr.Checkbox(value=False), gr.Number(value=0)
        control_status = gr.Textbox()
        gr.Button('Control API').click(control, [session, control_sequence, control_x, control_z,
                                                control_shift, control_slot], control_status,
                                        api_name='control', queue=False)
        inspect_status = gr.JSON()
        gr.Button('Session status').click(session_status, session, inspect_status,
                                           queue=False, api_name='session_status')

demo.queue(max_size=32)
if __name__ == '__main__':
    demo.launch(server_name='0.0.0.0', server_port=int(os.getenv('PORT', '7860')),
                show_error=True, footer_links=[])