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

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

def stream_output(process, q):
    for line in iter(process.stdout.readline, b''):
        q.put(line.decode('utf-8').strip())
    process.stdout.close()

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

    results = []
    manifest_logs = []

    # 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.", "Error: Cannot start manifest server without HF_TOKEN."

    manifest_process = None
    log_queue = queue.Queue()
    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.STDOUT, text=True, bufsize=1, universal_newlines=True)

        threading.Thread(target=stream_output, args=(manifest_process, log_queue), daemon=True).start()

        results.append("Started manifest server in background.")

        # Wait for the server to initialize (adjust time as needed)
        for _ in range(30):
            try:
                line = log_queue.get(timeout=1)
                manifest_logs.append(line)
                if "Running on" in line:  # Server is ready
                    break
            except queue.Empty:
                pass

        # 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 couldn't 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.")

        # Collect any remaining logs
        while True:
            try:
                line = log_queue.get_nowait()
                manifest_logs.append(line)
            except queue.Empty:
                break

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

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="Evaluation Output", lines=20)
    manifest_output = gr.Textbox(label="Manifest Server Logs", lines=20)

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

demo.launch()