import gradio as gr
import sqlite3
import requests
import os

def fetch_data(db_path, limit_links=None):
    if not os.path.exists(db_path):
        return [], []
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    cursor.execute('''CREATE TABLE IF NOT EXISTS links (
                        id INTEGER PRIMARY KEY,
                        url TEXT,
                        status TEXT DEFAULT 'pending'
                    )''')
    try:
        cursor.execute("SELECT api_key FROM accounts")
        api_keys = [row[0] for row in cursor.fetchall() if row[0]]
    except Exception as e:
        api_keys = []
        
    try:
        query = "SELECT id, url FROM links WHERE status='pending'"
        if limit_links:
            query += f" LIMIT {limit_links}"
            
        cursor.execute(query)
        link_records = cursor.fetchall()
        links = [row[1] for row in link_records if row[1]]
        
        if link_records:
            ids = [str(row[0]) for row in link_records]
            cursor.execute(f"UPDATE links SET status='processing' WHERE id IN ({','.join(ids)})")
            conn.commit()
    except Exception as e:
        links = []
        
    conn.close()
    return api_keys, links

def trigger_workers(db_path, worker_urls_str, links_per_worker):
    worker_urls = [u.strip() for u in worker_urls_str.split(',') if u.strip()]
    if not worker_urls:
        return "Error: No worker URLs provided."
        
    api_keys, links = fetch_data(db_path, limit_links=int(links_per_worker) * len(worker_urls))
    if not api_keys:
        return "Error: No API keys found in DB."
    if not links:
        return "Error: No pending links found in DB."
        
    link_chunks = [links[i:i + int(links_per_worker)] for i in range(0, len(links), int(links_per_worker))]
    
    logs = []
    logs.append(f"Loaded {len(api_keys)} API keys and {len(links)} links.")
    logs.append(f"Divided into {len(link_chunks)} worker chunks.")
    
    for idx, chunk in enumerate(link_chunks):
        if idx >= len(worker_urls):
            logs.append("More chunks than workers. Remaining chunks won't be sent.")
            break
            
        worker_url = worker_urls[idx]
        api_key = api_keys[idx % len(api_keys)]
        
        payload = {
            "links": chunk,
            "api_key": api_key
        }
        
        endpoint = f"{worker_url.rstrip('/')}/process"
        try:
            resp = requests.post(endpoint, json=payload, timeout=10)
            logs.append(f"Sent {len(chunk)} links to {worker_url}: {resp.status_code} - {resp.text}")
        except Exception as e:
            logs.append(f"Failed to send to {worker_url}: {str(e)}")
            
    return "\n".join(logs)

with gr.Blocks() as app:
    gr.Markdown("# YouTube Video Orchestrator Host")
    gr.Markdown("Deploy this to Hugging Face Spaces or run locally. It fetches pending links from your database and delegates them to FastAPI workers.")
    
    with gr.Row():
        db_input = gr.Textbox(label="DB Path", value="data/db/accounts.db")
        workers_input = gr.Textbox(label="Worker URLs (comma separated)", placeholder="http://localhost:7860, https://your-worker.hf.space")
        links_per_worker = gr.Number(label="Links per worker", value=2)
        
    run_btn = gr.Button("Run Orchestration", variant="primary")
    output_log = gr.Textbox(label="Logs", lines=10)
    
    run_btn.click(fn=trigger_workers, inputs=[db_input, workers_input, links_per_worker], outputs=output_log)

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)