import os
import gradio as gr
import threading
from flask import Flask, jsonify
from psycopg2 import connect
import psycopg2.extras

# Fetch environment variables
DB_NAME = os.getenv("DB_NAME")
DB_USER = os.getenv("DB_USER")
DB_PASSWORD = os.getenv("DB_PASSWORD")
DB_HOST = os.getenv("DB_HOST")
DB_PORT = os.getenv("DB_PORT")
APP_PASSWORD = os.getenv("APP_PASSWORD")

# Database connection
def get_db_connection():
    return connect(
        dbname=DB_NAME,
        user=DB_USER,
        password=DB_PASSWORD,
        host=DB_HOST,
        port=DB_PORT
    )

# Database operation functions
def get_db_stats():
    conn = get_db_connection()
    cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
    
    # Example: Get database size (this is a placeholder query, you can adapt it to your use case)
    cursor.execute("SELECT pg_size_pretty(pg_database_size(current_database())) AS size")
    db_stats = cursor.fetchone()
    
    conn.close()
    return db_stats['size']

def create_table(name):
    conn = get_db_connection()
    cursor = conn.cursor()

    # Create a basic table (customize the structure as needed)
    cursor.execute(f"""
    CREATE TABLE IF NOT EXISTS {name} (
        id SERIAL PRIMARY KEY,
        data TEXT
    );
    """)
    conn.commit()
    conn.close()
    return f"Table {name} created successfully!"

def delete_table(name):
    conn = get_db_connection()
    cursor = conn.cursor()

    # Delete table
    cursor.execute(f"DROP TABLE IF EXISTS {name};")
    conn.commit()
    conn.close()
    return f"Table {name} deleted successfully!"

# Flask app for HTTP access to the database
def create_http_server():
    app = Flask(__name__)

    @app.route("/db-stats", methods=["GET"])
    def db_stats():
        db_stats = get_db_stats()
        return jsonify({"db_size": db_stats})

    @app.route("/create-table/<name>", methods=["POST"])
    def create_table_route(name):
        result = create_table(name)
        return jsonify({"message": result})

    @app.route("/delete-table/<name>", methods=["DELETE"])
    def delete_table_route(name):
        result = delete_table(name)
        return jsonify({"message": result})

    return app

# Manage server and database operations in the Gradio UI
def manage_servers(password, action, name=None, port=None):
    if password != APP_PASSWORD:
        return "Invalid password!"

    if action == "start":
        return start_server(name, port)
    elif action == "stop":
        return stop_server(name)
    elif action == "db-stats":
        return get_db_stats()
    elif action == "create-table" and name:
        return create_table(name)
    elif action == "delete-table" and name:
        return delete_table(name)
    else:
        return "Invalid action!"

# Start and stop server functions (same as before)
server_threads = {}

def create_server(name, port):
    app = Flask(name)

    @app.route("/")
    def home():
        return f"Server {name} is running on port {port}!"

    return app

def start_server(name, port):
    if name in server_threads:
        return f"Server {name} is already running!"

    app = create_server(name, port)
    
    def run():
        app.run(host="0.0.0.0", port=port, debug=False, use_reloader=False)

    thread = threading.Thread(target=run, daemon=True)
    thread.start()
    server_threads[name] = thread
    return f"Server {name} started on port {port}!"

def stop_server(name):
    if name not in server_threads:
        return f"Server {name} is not running!"

    server_threads.pop(name)
    return f"Server {name} stopped!"

# Gradio UI
with gr.Blocks() as server_manager:
    gr.Markdown("# Server Manager")
    
    with gr.Row():
        password = gr.Textbox(label="Password", type="password")
    
    with gr.Row():
        action = gr.Radio(["start", "stop", "db-stats", "create-table", "delete-table"], label="Action")
        name = gr.Textbox(label="Server or Table Name")  # Removed optional=True
        port = gr.Number(label="Port", visible=False)  # Making port field visible only when starting a server
    
    with gr.Row():
        submit = gr.Button("Submit")
    
    output = gr.Textbox(label="Output")
    
    submit.click(
        manage_servers,
        inputs=[password, action, name, port],
        outputs=output
    )

# Start HTTP server for database operations
def start_http_server():
    http_server = create_http_server()
    http_thread = threading.Thread(target=http_server.run, kwargs={"host": "0.0.0.0", "port": 5000}, daemon=True)
    http_thread.start()

# Start everything
if __name__ == "__main__":
    start_http_server()
    server_manager.launch()