import gradio as gr
import sqlite3
import requests
import os

def fetch_data(db_path, limit_links=None):
    os.makedirs(os.path.dirname(db_path), exist_ok=True)
    if not os.path.exists(db_path):
        # We can either create it or return empty. I will let sqlite create it if it doesn't exist.
        pass
    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)

import json
import uuid
import shutil
import subprocess
from pathlib import Path

def get_kaggle_username():
    kaggle_json = Path.home() / ".kaggle" / "kaggle.json"
    if not kaggle_json.exists():
        return "your-username"
    with open(kaggle_json, "r") as f:
        creds = json.load(f)
        return creds.get("username", "your-username")

def trigger_kaggle_test():
    username = get_kaggle_username()
    worker_title = f"yt-test-worker-{uuid.uuid4().hex[:6]}"
    worker_dir = Path("kaggle_test_worker")
    
    if worker_dir.exists():
        shutil.rmtree(worker_dir)
    worker_dir.mkdir(parents=True)
    
    notebook_content = {
        "cells": [
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "print('Hello from Kaggle Worker!')\n",
                    "print('Test successful.')\n"
                ]
            }
        ],
        "metadata": {
            "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
            "language_info": {"name": "python"}
        },
        "nbformat": 4,
        "nbformat_minor": 4
    }
    
    with open(worker_dir / "notebook.ipynb", "w") as f:
        json.dump(notebook_content, f, indent=2)
        
    metadata = {
        "id": f"{username}/{worker_title}",
        "title": worker_title,
        "code_file": "notebook.ipynb",
        "language": "python",
        "kernel_type": "notebook",
        "is_private": True,
        "enable_gpu": False,
        "enable_internet": True,
        "dataset_sources": [],
        "competition_sources": [],
        "kernel_sources": []
    }
    
    with open(worker_dir / "kernel-metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)
        
    logs = [f"Prepared Kaggle Test Worker: {worker_title}"]
    
    # Use kaggle from PATH (assuming it is installed via requirements.txt)
    command = ["kaggle", "kernels", "push", "-p", str(worker_dir)]
    logs.append(f"Pushing {worker_title} to Kaggle...")
    
    try:
        res = subprocess.run(command, capture_output=True, text=True)
        if res.returncode == 0:
            logs.append("Push successful.")
            logs.append(res.stdout)
        else:
            logs.append("Push failed.")
            logs.append(res.stderr)
    except Exception as e:
        logs.append(f"Error running kaggle command: {e}")
        logs.append("Make sure Kaggle API credentials are provided in ~/.kaggle/kaggle.json")
        
    return "\n".join(logs)

def create_worker_notebook(worker_dir, config_name):
    notebook_content = {
        "cells": [
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "!pip install openai huggingface_hub yt-dlp opencv-python-headless\n"
                ]
            },
            {
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "outputs": [],
                "source": [
                    "import json\n",
                    "import os\n",
                    "import subprocess\n",
                    "import sys\n",
                    "\n",
                    f"with open('{config_name}') as f:\n",
                    "    config = json.load(f)\n",
                    "\n",
                    "links = config['links']\n",
                    "api_key = config['api_key']\n",
                    "os.environ['FEATHERLESS_API_KEY'] = api_key\n",
                    "\n",
                    "print(f'Starting worker with {len(links)} links using API Key: {api_key[:5]}...')\n",
                    "\n",
                    "for link in links:\n",
                    "    print(f'\\n[+] Processing: {link}')\n",
                    "    # 1. Download Video\n",
                    "    dl_cmd = [sys.executable, 'worker/youtube_worker/downloader.py', link, '--resolution', 'HD', '-o', 'test_output']\n",
                    "    res = subprocess.run(dl_cmd, capture_output=True, text=True)\n",
                    "    if res.returncode != 0:\n",
                    "        print(f'Download failed:\\n{res.stderr}')\n",
                    "        continue\n",
                    "    \n",
                    "    # Find the downloaded video path from stdout (last line)\n",
                    "    lines = res.stdout.strip().split('\\n')\n",
                    "    video_path = lines[-1]\n",
                    "    if not os.path.exists(video_path):\n",
                    "        print(f'Could not find downloaded video at: {video_path}')\n",
                    "        continue\n",
                    "        \n",
                    "    print(f'[+] Video downloaded to: {video_path}')\n",
                    "    \n",
                    "    # 2. Curate Model (Video Worker)\n",
                    "    prompt = 'Analyze this video comprehensively. Focus on visuals, objects, and text.'\n",
                    "    vw_cmd = [sys.executable, 'worker/models_worker/video_worker.py', '--video-url', video_path, '--prompt', prompt, '--api-key', api_key]\n",
                    "    print('[+] Running model curation...')\n",
                    "    subprocess.run(vw_cmd)\n",
                    "    \n",
                    "    # 3. Optional: Upload to HF\n",
                    "    # hf_cmd = [sys.executable, 'worker/hf_uploader.py', 'test_output']\n",
                    "    # subprocess.run(hf_cmd)\n",
                    "\n",
                    "print('\\nAll tasks completed!')\n"
                ]
            }
        ],
        "metadata": {
            "kernelspec": {
                "display_name": "Python 3",
                "language": "python",
                "name": "python3"
            },
            "language_info": {
                "name": "python"
            }
        },
        "nbformat": 4,
        "nbformat_minor": 4
    }
    
    with open(worker_dir / "notebook.ipynb", "w") as f:
        json.dump(notebook_content, f, indent=2)

def trigger_kaggle_workers(db_path, links_per_worker, max_workers):
    api_keys, links = fetch_data(db_path, limit_links=int(links_per_worker) * int(max_workers))
    if not api_keys:
        return "Error: No API keys found in DB."
    if not links:
        return "Error: No pending links found in DB."
        
    username = get_kaggle_username()
    workers_dir = Path("kaggle_workers_out")
    if workers_dir.exists():
        shutil.rmtree(workers_dir)
    workers_dir.mkdir(parents=True)
    
    logs = []
    logs.append(f"Loaded {len(api_keys)} API keys and {len(links)} links.")
    
    link_chunks = [links[i:i + int(links_per_worker)] for i in range(0, len(links), int(links_per_worker))]
    logs.append(f"Divided into {len(link_chunks)} worker chunks.")
    
    for idx, chunk in enumerate(link_chunks):
        worker_id = f"worker-{uuid.uuid4().hex[:8]}"
        worker_title = f"yt-worker-{idx+1}-{uuid.uuid4().hex[:4]}"
        
        api_key = api_keys[idx % len(api_keys)]
        worker_path = workers_dir / worker_id
        worker_path.mkdir(parents=True)
        
        # Copy worker if exists locally
        if Path("worker").exists():
            shutil.copytree("worker", worker_path / "worker")
        else:
            # We copy from parent directory if this is running locally, otherwise it's an error on HF unless packed
            parent_worker = Path("../worker")
            if parent_worker.exists():
                shutil.copytree(parent_worker, worker_path / "worker")
            else:
                logs.append("Warning: 'worker' directory not found. Kaggle notebook may fail.")
                
        if Path("cookies.txt").exists():
            shutil.copy("cookies.txt", worker_path / "cookies.txt")
        elif Path("../cookies.txt").exists():
            shutil.copy("../cookies.txt", worker_path / "cookies.txt")
            
        if Path("cookies.json").exists():
            shutil.copy("cookies.json", worker_path / "cookies.json")
        elif Path("../cookies.json").exists():
            shutil.copy("../cookies.json", worker_path / "cookies.json")
            
        config_name = "worker_config.json"
        config_data = {
            "links": chunk,
            "api_key": api_key,
        }
        with open(worker_path / config_name, "w") as f:
            json.dump(config_data, f, indent=2)
            
        create_worker_notebook(worker_path, config_name)
        
        metadata = {
            "id": f"{username}/{worker_title}",
            "title": worker_title,
            "code_file": "notebook.ipynb",
            "language": "python",
            "kernel_type": "notebook",
            "is_private": True,
            "enable_gpu": False,
            "enable_internet": True,
            "dataset_sources": [],
            "competition_sources": [],
            "kernel_sources": []
        }
        with open(worker_path / "kernel-metadata.json", "w") as f:
            json.dump(metadata, f, indent=2)
            
        logs.append(f"Prepared Worker {idx+1}: {worker_title} with {len(chunk)} links.")
        
        command = ["kaggle", "kernels", "push", "-p", str(worker_path)]
        logs.append(f"Pushing {worker_title} to Kaggle...")
        try:
            res = subprocess.run(command, capture_output=True, text=True)
            if res.returncode == 0:
                logs.append("Push successful.")
            else:
                logs.append("Push failed.")
                logs.append(res.stderr)
        except Exception as e:
            logs.append(f"Error: {e}")
            
    return "\n".join(logs)

def sync_db(db_path, api_keys_text, links_text):
    if not db_path:
        return "Error: DB Path is empty."
    
    os.makedirs(os.path.dirname(db_path), exist_ok=True)
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # Ensure tables exist
    cursor.execute('''CREATE TABLE IF NOT EXISTS links (
                        id INTEGER PRIMARY KEY,
                        url TEXT,
                        status TEXT DEFAULT 'pending'
                    )''')
    
    cursor.execute('''CREATE TABLE IF NOT EXISTS accounts (
                        id INTEGER PRIMARY KEY,
                        email TEXT,
                        api_key TEXT,
                        created_at TEXT DEFAULT CURRENT_TIMESTAMP
                    )''')
    
    logs = []
    
    # Add API keys
    api_keys = [k.strip() for k in api_keys_text.split('\n') if k.strip()]
    if api_keys:
        added_keys = 0
        for key in api_keys:
            try:
                # Basic check to avoid exact duplicates
                cursor.execute("SELECT id FROM accounts WHERE api_key=?", (key,))
                if not cursor.fetchone():
                    cursor.execute("INSERT INTO accounts (email, api_key) VALUES (?, ?)", (f"user_{uuid.uuid4().hex[:4]}@example.com", key))
                    added_keys += 1
            except Exception as e:
                logs.append(f"Error adding key {key}: {e}")
        conn.commit()
        logs.append(f"Successfully added {added_keys} new API keys.")
        
    # Add Links
    links = [l.strip() for l in links_text.split('\n') if l.strip()]
    if links:
        added_links = 0
        for link in links:
            try:
                # Basic check to avoid exact duplicates
                cursor.execute("SELECT id FROM links WHERE url=?", (link,))
                if not cursor.fetchone():
                    cursor.execute("INSERT INTO links (url, status) VALUES (?, 'pending')", (link,))
                    added_links += 1
            except Exception as e:
                logs.append(f"Error adding link {link}: {e}")
        conn.commit()
        logs.append(f"Successfully added {added_links} new pending links.")
        
    conn.close()
    
    if not api_keys and not links:
        return "No API keys or links provided to sync."
        
    return "\n".join(logs)

def run_v3_launch(max_workers, videos_per_worker):
    cmd = ["python", "pipeline/orchestrator.py", "launch", 
           "--max-workers", str(int(max_workers)), 
           "--videos-per-worker", str(int(videos_per_worker)), 
           "--push"]
    try:
        # Run in a subprocess and capture output
        res = subprocess.run(cmd, capture_output=True, text=True, cwd=os.path.dirname(os.path.abspath(__file__)))
        return f"--- LAUNCH OUTPUT ---\n{res.stdout}\n{res.stderr}"
    except Exception as e:
        return f"Error: {e}"

def run_v3_sync():
    cmd = ["python", "pipeline/orchestrator.py", "sync"]
    try:
        res = subprocess.run(cmd, capture_output=True, text=True, cwd=os.path.dirname(os.path.abspath(__file__)))
        return f"--- SYNC OUTPUT ---\n{res.stdout}\n{res.stderr}"
    except Exception as e:
        return f"Error: {e}"

def run_v3_status():
    cmd = ["python", "pipeline/orchestrator.py", "status"]
    try:
        res = subprocess.run(cmd, capture_output=True, text=True, cwd=os.path.dirname(os.path.abspath(__file__)))
        return f"--- STATUS OUTPUT ---\n{res.stdout}\n{res.stderr}"
    except Exception as e:
        return f"Error: {e}"

def run_v3_retry():
    cmd = ["python", "pipeline/orchestrator.py", "retry"]
    try:
        res = subprocess.run(cmd, capture_output=True, text=True, cwd=os.path.dirname(os.path.abspath(__file__)))
        return f"--- RETRY OUTPUT ---\n{res.stdout}\n{res.stderr}"
    except Exception as e:
        return f"Error: {e}"

def run_verify_repo(repo_id):
    cmd = ["python", "check_repo_db.py", "--repo-id", repo_id]
    try:
        res = subprocess.run(cmd, capture_output=True, text=True, cwd=os.path.dirname(os.path.abspath(__file__)))
        return f"--- REPO VERIFICATION OUTPUT ---\n{res.stdout}\n{res.stderr}"
    except Exception as e:
        return f"Error: {e}"

autopilot_process = None

def toggle_autopilot(videos_per_worker):
    global autopilot_process
    
    if autopilot_process is not None and autopilot_process.poll() is None:
        autopilot_process.terminate()
        autopilot_process = None
        return "🛑 Autopilot stopped."
    else:
        cmd = ["python", "pipeline/orchestrator.py", "run", "--videos-per-worker", str(int(videos_per_worker))]
        autopilot_process = subprocess.Popen(cmd, cwd=os.path.dirname(os.path.abspath(__file__)))
        return "🚀 Autopilot started! It is now running continuously in the background, checking Kaggle every 2 minutes and instantly launching new workers whenever slots free up."

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 or Kaggle.")
    
    with gr.Tabs():
        with gr.Tab("V3 Pipeline Orchestrator"):
            gr.Markdown("Directly run the advanced `pipeline/orchestrator.py` commands from this dashboard!")
            
            with gr.Row():
                v3_max_workers = gr.Number(label="Max Workers", value=5)
                v3_videos = gr.Number(label="Videos per Worker", value=250)
                
            with gr.Row():
                v3_launch_btn = gr.Button("Launch Workers (--push)", variant="primary")
                v3_autopilot_btn = gr.Button("Toggle Autopilot", variant="secondary")
                v3_sync_btn = gr.Button("Sync DB & Status")
                v3_status_btn = gr.Button("Check Status")
                v3_retry_btn = gr.Button("Retry Failed Videos")
                
            gr.Markdown("---")
            gr.Markdown("### Verify Database against any HF Repository")
            gr.Markdown("List files in a target HF repo and mark any matching pending videos as 'done' in your tracker database.")
            with gr.Row():
                verify_repo_id = gr.Textbox(label="Hugging Face Repo ID", value="AdhyanshVerma/YT")
                verify_btn = gr.Button("Verify and Update DB", variant="primary")
                
            v3_output_log = gr.Textbox(label="Terminal Output", lines=15)
            
            v3_launch_btn.click(fn=run_v3_launch, inputs=[v3_max_workers, v3_videos], outputs=v3_output_log)
            v3_autopilot_btn.click(fn=toggle_autopilot, inputs=[v3_videos], outputs=v3_output_log)
            v3_sync_btn.click(fn=run_v3_sync, inputs=[], outputs=v3_output_log)
            v3_status_btn.click(fn=run_v3_status, inputs=[], outputs=v3_output_log)
            v3_retry_btn.click(fn=run_v3_retry, inputs=[], outputs=v3_output_log)
            verify_btn.click(fn=run_verify_repo, inputs=[verify_repo_id], outputs=v3_output_log)

        with gr.Tab("Database Sync"):
            gr.Markdown("Sync (Add) new API keys and YouTube Links into your database.")
            with gr.Row():
                sync_db_input = gr.Textbox(label="DB Path", value="/data/accounts.db")
            with gr.Row():
                api_keys_input = gr.Textbox(label="API Keys (One per line)", lines=5)
                links_input = gr.Textbox(label="YouTube Links (One per line)", lines=5)
                
            sync_btn = gr.Button("Sync to Database", variant="primary")
            sync_log = gr.Textbox(label="Sync Logs", lines=5)
            
            sync_btn.click(fn=sync_db, inputs=[sync_db_input, api_keys_input, links_input], outputs=sync_log)

        with gr.Tab("FastAPI Workers"):
            with gr.Row():
                db_input = gr.Textbox(label="DB Path", value="/data/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)

        with gr.Tab("Simple Kaggle Workers"):
            gr.Markdown("Trigger simple Kaggle notebook workers. Note: You must mount `kaggle.json` to `~/.kaggle/kaggle.json` in the container for this to work.")
            
            with gr.Row():
                k_db_input = gr.Textbox(label="DB Path", value="/data/accounts.db")
                k_links_per_worker = gr.Number(label="Links per worker", value=2)
                k_max_workers = gr.Number(label="Max Workers", value=5)
                
            k_run_btn = gr.Button("Run Simple Kaggle Orchestration", variant="primary")
            k_test_run_btn = gr.Button("Run Simple Kaggle Test Worker", variant="secondary")
            kaggle_log = gr.Textbox(label="Kaggle Logs", lines=15)
            
            k_run_btn.click(fn=trigger_kaggle_workers, inputs=[k_db_input, k_links_per_worker, k_max_workers], outputs=kaggle_log)
            k_test_run_btn.click(fn=trigger_kaggle_test, inputs=[], outputs=kaggle_log)

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