import gradio as gr
import json
import time

status_file = "status.json"
command_file = "command.json"
output_file = "output.json"

HEARTBEAT_TIMEOUT = 20  # seconds

# Initialize storage files
with open(status_file, "w") as f:
    json.dump({}, f)
with open(command_file, "w") as f:
    json.dump({"command": ""}, f)
with open(output_file, "w") as f:
    json.dump({"output": ""}, f)

# Register PC as online
def register_pc(name, password):
    with open(status_file, "r") as f:
        status = json.load(f)
    status[name] = {"password": password, "timestamp": time.time()}
    with open(status_file, "w") as f:
        json.dump(status, f)
    return f"PC '{name}' registered as online."

# Check if PC is online
def check_pc(name):
    with open(status_file, "r") as f:
        status = json.load(f)
    if name in status and time.time() - status[name]["timestamp"] < HEARTBEAT_TIMEOUT:
        return "Online"
    return "Offline"

# Send command (or register if command is "REGISTER")
def send_command(name, password, command):
    if command == "REGISTER":
        return register_pc(name, password)
    with open(status_file, "r") as f:
        status = json.load(f)
    if name in status and status[name]["password"] == password:
        with open(command_file, "w") as f:
            json.dump({"command": command}, f)
        return f"Command sent to {name}."
    return "Authentication failed."

# Get command (for PC side polling)
def get_command():
    with open(command_file, "r") as f:
        cmd_data = json.load(f)
    return cmd_data.get("command", "")

# Upload command output from PC with timestamps
def upload_output(name, password, output, sent_time):
    with open(status_file, "r") as f:
        status = json.load(f)
    if name in status and status[name]["password"] == password:
        received_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        # Format the output with "Time Sent", a divider, then "Time Received", another divider, and the actual output.
        combined_output = (
            f"Time Sent: {sent_time}\n"
            "--------------------\n"
            f"Time Received: {received_time}\n"
            "--------------------\n"
            f"{output}"
        )
        with open(output_file, "w") as f:
            json.dump({"output": combined_output}, f)
        return "Output uploaded successfully."
    return "Authentication failed."

# Retrieve command output (to view on Hugging Face)
def get_output():
    with open(output_file, "r") as f:
        data = json.load(f)
    return data.get("output", "")

with gr.Blocks() as ui:
    # PC info and status
    with gr.Row():
        name_input = gr.Textbox(label="PC Name")
        pass_input = gr.Textbox(label="Password", type="password")
    check_btn = gr.Button("Check Online Status")
    check_output = gr.Textbox(label="Status")
    check_btn.click(check_pc, inputs=name_input, outputs=check_output)

    # Command sending row
    with gr.Row():
        cmd_input = gr.Textbox(label="Command", lines=5)  # Multi-line for larger scripts
        send_btn = gr.Button("Send Command")
        send_output = gr.Textbox(label="Response", lines=2)
    send_btn.click(send_command, inputs=[name_input, pass_input, cmd_input], outputs=send_output)
    
    # Visible output display row
    with gr.Row():
        get_output_btn = gr.Button("Get Command Output")
        output_display = gr.Textbox(label="Command Output", lines=10)
    get_output_btn.click(get_output, inputs=[], outputs=output_display)
    
    # Hidden endpoints for PC polling and uploading output:
    hidden_get_cmd_btn = gr.Button("Hidden Get Command", visible=False)
    hidden_cmd_output = gr.Textbox(visible=False)
    hidden_get_cmd_btn.click(get_command, inputs=[], outputs=hidden_cmd_output, api_name="/get_command")
    
    # Hidden components for uploading output: includes a textbox for sent_time.
    output_input = gr.Textbox(label="Output", visible=False)
    sent_time_input = gr.Textbox(label="Sent Time", visible=False)
    hidden_upload_btn = gr.Button("Hidden Upload Output", visible=False)
    hidden_upload_output = gr.Textbox(visible=False)
    hidden_upload_btn.click(
        upload_output, 
        inputs=[name_input, pass_input, output_input, sent_time_input],
        outputs=hidden_upload_output, 
        api_name="/upload_output"
    )

ui.launch(share=True)