import gradio as gr
import subprocess
import spaces
import torch
import os
import re

zero = torch.Tensor([0]).cuda()
print(zero.device)  # <-- 'cpu' 🤔

@spaces.GPU
def run_evaluation(model_name):
    print(zero.device)  # <-- 'cuda:0' 🤗

    results = []

    # Use the secret HF token from the Hugging Face space
    if "HF_TOKEN" not in os.environ:
        return "Error: HF_TOKEN not found in environment variables."

    manifest_process = None
    try:
        # Start manifest server in background with explicit CUDA_VISIBLE_DEVICES
        manifest_cmd = f"""
        cd duckdb-nsql/ && 
        CUDA_VISIBLE_DEVICES=0 HF_TOKEN={os.environ['HF_TOKEN']}  python -m manifest.api.app \
        --model_type huggingface \
        --model_generation_type text-generation \
        --model_name_or_path {model_name} \
        --fp16 \
        --device 0
        """
        manifest_process = subprocess.Popen(manifest_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        results.append("Started manifest server in background.")

        # Run inference
        inference_cmd = f"""
        cd duckdb-nsql/ && 
        python eval/predict.py \
        predict \
        eval/data/dev.json \
        eval/data/tables.json \
        --output-dir output/ \
        --stop-tokens ';' \
        --overwrite-manifest \
        --manifest-client huggingface \
        --manifest-connection http://localhost:5000 \
        --prompt-format duckdbinstgraniteshort
        """
        inference_result = subprocess.run(inference_cmd, shell=True, check=True, capture_output=True, text=True)
        results.append("Inference completed.")

        # Extract JSON file path from inference output
        json_path_match = re.search(r'(.*\.json)', inference_result.stdout)
        if not json_path_match:
            raise ValueError("Could not find JSON file path in inference output")
        json_file = os.path.basename(json_path_match.group(1))
        results.append(f"Generated JSON file: {json_file}")

        # Run evaluation
        eval_cmd = f"""
        cd duckdb-nsql/ && 
        python eval/evaluate.py evaluate \
        --gold eval/data/dev.json \
        --db eval/data/databases/ \
        --tables eval/data/tables.json \
        --output-dir output/ \
        --pred output/{json_file}
        """
        eval_result = subprocess.run(eval_cmd, shell=True, check=True, capture_output=True, text=True)

        # Extract and format metrics from eval output
        metrics = eval_result.stdout
        if metrics:
            results.append(f"Evaluation completed:\n{metrics}")
        else:
            results.append("Evaluation completed, but get metrics.")

    except subprocess.CalledProcessError as e:
        results.append(f"Error occurred: {str(e)}")
        results.append(f"Command output: {e.output}")
    except Exception as e:
        results.append(f"An unexpected error occurred: {str(e)}")
    finally:
        # Terminate the background manifest server
        if manifest_process:
            manifest_process.terminate()
            results.append("Terminated manifest server.")

    return "\n\n".join(results)

with gr.Blocks() as demo:
    gr.Markdown("# DuckDB SQL Evaluation App")

    model_name = gr.Textbox(label="Model Name (e.g., Qwen/Qwen2.5-7B-Instruct)")
    start_btn = gr.Button("Start Evaluation")
    output = gr.Textbox(label="Output", lines=20)

    start_btn.click(fn=run_evaluation, inputs=[model_name], outputs=output)

demo.launch()