import os

# Import spaces BEFORE torch/transformers when using ZeroGPU.
MODE = os.environ.get('EXECUTION_MODE', 'zerogpu').lower()
if MODE not in {'zerogpu','cpu','cuda'}:
    raise ValueError('EXECUTION_MODE must be zerogpu, cpu or cuda')
if MODE == 'zerogpu':
    import spaces

import json
import logging
import gradio as gr
from runtime import load_runtime, parse_criteria, prepare_inputs

logging.basicConfig(level=logging.INFO)
LOGGER = logging.getLogger('sroute')
RUNTIME = None


def run_model(encoded, criteria, lengths):
    return RUNTIME.predict(encoded, criteria, lengths)


if MODE == 'zerogpu':
    # CUDA model placement occurs during startup below, not inside this function.
    run_model = spaces.GPU(duration=30)(run_model)


def predict(state, question, criteria_json):
    try:
        criteria = parse_criteria(criteria_json)
        encoded, lengths = prepare_inputs(RUNTIME.tokenizer,state,question,criteria)
        return run_model(encoded,criteria,lengths)
    except ValueError as exc:
        raise gr.Error(str(exc)) from exc
    except Exception as exc:
        LOGGER.exception('Prediction failed')
        raise gr.Error('Inference failed. Check the Space Runtime logs for the full traceback.') from exc


def build_demo():
    criteria = json.dumps({
        'opt_1':'Customer Service - general business questions and routine account updates',
        'opt_2':'Technical Support - login failures, password resets, application errors',
        'opt_3':'Fraud Prevention - unauthorized transactions and suspected fraud'},indent=2)
    return gr.Interface(fn=predict,
        inputs=[gr.Textbox(label='State / customer message',lines=4),
                gr.Textbox(label='Routing question',value='Which specific department or action should handle this request?'),
                gr.Textbox(label='Routing criteria: JSON object, 2 to 4 options',value=criteria,lines=7)],
        outputs=gr.JSON(label='Decision'),
        title='sroute — Business Routing POC',
        description='Routing, with experimental escalation and severity estimates. '
                    'Auxiliary outputs depend on the first routing option in this checkpoint. '
                    'Inputs above 256 tokens are rejected instead of silently truncated.',
        examples=[['I cannot log in to my account. Please help me reset my password.',
                   'Which specific department or action should handle this request?',criteria]],
        cache_examples=False)


if __name__ == '__main__':
    # ZeroGPU requires cuda placement during startup even though allocation is dynamic.
    device = 'cuda' if MODE in {'zerogpu','cuda'} else 'cpu'
    RUNTIME = load_runtime(device=device)
    LOGGER.info('Loaded checkpoint revision=%s mode=%s',RUNTIME.metadata['model_revision'],MODE)
    demo = build_demo()
    demo.queue(default_concurrency_limit=1,max_size=16)
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", "7860")),
        ssr_mode=False,
        show_error=False,
    )