import gradio as gr

from urllib.parse import urlparse

import subprocess
import threading

import sys
import os


LOG_FILE = "output.log"

def read_logs():
    sys.stdout.flush()
    try:
        with open(LOG_FILE, "r") as f:
            return f.read()
    except Exception:
        return "ML worker not running"


previous_url = ""
ml_worker = None

def read_status():
    if ml_worker:
        return f"ML worker serving {previous_url}"
    elif len(previous_url):
        return f"ML worker exited for {previous_url}"
    else:
        return "ML worker not started"


def run_ml_worker(url, api_key, hf_token):
    global ml_worker, previous_url
    previous_url = url
    ml_worker = subprocess.Popen(
        [
            "giskard", "worker", "start",
            "-u", f"{url}", "-k", f"{api_key}", "-t", f"{hf_token}"
        ],
        stdout=open(LOG_FILE, "w"), stderr=subprocess.STDOUT
    )
    args = ml_worker.args
    print(f"Process {args} exited with {ml_worker.wait()}")
    ml_worker = None


def stop_ml_worker():
    global ml_worker, previous_url
    if ml_worker is not None:
        print(f"Stopping ML worker for {previous_url}")
        ml_worker.terminate()
        ml_worker = None
        print("ML worker stopped")
        return "ML worker stopped"
    return "ML worker not started"


def start_ml_worker(url, api_key, hf_token):
    if not url or len(url) < 1:
        return "Please provide URL of Giskard"

    # Always run an external ML worker
    stop_ml_worker()

    print(f"Starting ML worker for {url}")
    thread = threading.Thread(target=run_ml_worker, args=(url, api_key, hf_token))
    thread.start()
    return f"ML worker running for {url}"

theme = gr.themes.Soft(
    primary_hue="green",
)

with gr.Blocks(theme=theme) as iface:
    with gr.Row():
        with gr.Column():
            url_input = gr.Textbox(label="Giskard Hub URL")
            api_key_input = gr.Textbox(
                label="Giskard Hub API Key",
                value=os.environ.get("GSK_API_KEY"),
                placeholder="gsk-xxxxxxxxxxxxxxxxxxxxxxxxxxxx",
            )
            hf_token_input = gr.Textbox(label="Hugging Face Spaces Token")

        output = gr.Textbox(label="Status")
    with gr.Row():
        run_btn = gr.Button("Run", variant="primary")
        run_btn.click(start_ml_worker, [url_input, api_key_input, hf_token_input], output)

        stop_btn = gr.Button("Stop", variant="stop")
        stop_btn.click(stop_ml_worker, None, output)

    logs = gr.Textbox(label="Giskard ML worker log:")
    iface.load(read_logs, None, logs, every=0.5)
    iface.load(read_status, None, output, every=5)

iface.queue()
iface.launch()