app.py

import gradio as gr
import time
import os
from backend import (
    init_db,
    create_user,
    get_user_by_username,
    verify_password,
    create_project,
    get_user_projects,
    get_project,
    queue_job,
)
import psutil
import torch

# --- UI HELPER FUNCTIONS ---

def get_system_stats():
    """Returns a markdown string with current RAM and VRAM usage."""
    ram_used = round(psutil.virtual_memory().used / (1024**3), 2)
    ram_total = round(psutil.virtual_memory().total / (1024**3), 2)
    stats_str = f"📊 System RAM: {ram_used}GB / {ram_total}GB"
    if torch.cuda.is_available():
        vram_reserved = round(torch.cuda.memory_reserved() / (1024**3), 2)
        vram_total = round(torch.cuda.get_device_properties(0).total_memory / (1024**3), 2)
        stats_str += f" | VRAM: {vram_reserved}GB / {vram_total}GB"
    return stats_str

def format_projects_list(projects):
    """Formats a list of project records into a readable markdown string."""
    if not projects:
        return "No projects yet. Start one on the left!"
    
    project_lines = []
    for p in projects:
        status_icon = {"queued": "🕒", "running": "⚙️", "completed": "✅", "failed": "❌"}.get(p['status'], "❓")
        project_lines.append(f"{status_icon} **[{p['id']}] {p['title']}** ({p['status']})")
    return "\n".join(project_lines)

def get_project_choices(user_id):
    """Fetches user projects and formats them for a dropdown."""
    if not user_id: return []
    projects = get_user_projects(user_id)
    return [(f"[{p['id']}] {p['title']}", p['id']) for p in projects] if projects else []

def format_files_list(projects):
    """Formats completed projects into downloadable file links."""
    completed_projects = [p for p in projects if p['status'] == 'completed']
    if not completed_projects:
        return "No completed projects yet. Your finished work will appear here!"
    
    file_lines = []
    for p in completed_projects:
        zip_path = p.get('zip_path')
        if zip_path and os.path.exists(zip_path):
            file_lines.append(f"📥 **[{p['id']}] {p['title']}** - [Download ZIP]({zip_path})")
        else:
            file_lines.append(f"⚠️ **[{p['id']}] {p['title']}** - ZIP file missing")
    return "\n".join(file_lines)

# --- CORE LOGIC HANDLERS FOR UI ---

def login_and_setup(username, password, user_state):
    """Handles login/signup and prepares the main UI if successful."""
    user = get_user_by_username(username)
    user_id, welcome_message = None, ""

    if user and verify_password(password, user['password_hash']):
        user_id = user['id']
        welcome_message = f"Welcome back, {username}! 👨‍💻"
    else:
        new_user_id = create_user(username, password)
        if new_user_id:
            user_id = new_user_id
            welcome_message = f"Account created! Welcome, {username}!"
        else:
            return user_state, gr.update(), gr.update(value="Login failed or username is taken.")

    if user_id:
        user_state['user_id'] = user_id
        user_state['username'] = username
        projects = get_user_projects(user_id)
        
        return (
            user_state,
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(value=welcome_message),
            gr.update(value=format_projects_list(projects)),
            gr.update(choices=get_project_choices(user_id), value=None),
            gr.update(value=format_files_list(projects))
        )
    return user_state, gr.update(), gr.update()


def start_project_flow(prompt, user_state):
    """The complete, seamless flow for starting a new project."""
    user_id = user_state.get('user_id')
    if not user_id:
        return "Error: Not logged in.", "", "", gr.update(value=""), gr.update(), gr.update(), gr.update()

    project_id = create_project(user_id, "New AI Project", prompt)
    if not project_id:
        return "Error: Could not create project.", "", "", gr.update(value=""), gr.update(), gr.update(), gr.update()
        
    queue_job(project_id, user_id, prompt)
    
    projects = get_user_projects(user_id)
    return (
        f"✅ Project #{project_id} started! Switching to logs...",
        gr.update(value=""),
        gr.update(value=format_projects_list(projects)),
        gr.update(choices=get_project_choices(user_id), value=project_id),
        gr.Tabs(selected="logs_tab"),
        gr.update(value=format_files_list(projects))
    )

def refresh_all_projects(user_state):
    """Manually refresh project lists on both tabs."""
    user_id = user_state.get('user_id')
    if not user_id: 
        return gr.update(), gr.update(), gr.update()
    
    projects = get_user_projects(user_id)
    return (
        gr.update(value=format_projects_list(projects)),
        gr.update(choices=get_project_choices(user_id)),
        gr.update(value=format_files_list(projects))
    )

# NEW: Refresh files tab button
def refresh_files(user_state):
    user_id = user_state.get('user_id')
    if not user_id: return gr.update()
    projects = get_user_projects(user_id)
    return gr.update(value=format_files_list(projects))

# --- CUSTOM GENERATOR LOOPS FOR LIVE UPDATES ---

def system_stats_loop():
    """A generator that yields system stats every 5 seconds, forever."""
    while True:
        yield get_system_stats()
        time.sleep(5)

def stream_logs(project_id):
    """A generator that streams logs and provides a download link when complete."""
    if not project_id:
        yield "⬅️ Select a project from the dropdown to see its live logs.", gr.update(visible=False)
        return

    yield "Fetching logs...", gr.update(visible=False)

    while True:
        project = get_project(project_id)
        if project:
            logs = project['logs'] if project['logs'] else f"🕒 Project #{project_id} is queued or starting... Logs will appear here shortly."
            
            if project['status'] == 'completed':
                final_log = logs + "\n\n✅ **Project Completed!** Your file is ready for download."
                zip_path = project['zip_path']
                if zip_path and os.path.exists(zip_path):
                    yield final_log, gr.update(value=zip_path, visible=True)
                else:
                    yield logs + "\n\n⚠️ **Warning:** Project completed, but the ZIP file was not found.", gr.update(visible=False)
                break
            
            elif project['status'] == 'failed':
                yield logs, gr.update(visible=False)
                break
            
            yield logs, gr.update(visible=False)
        else:
            yield f"Error: Project #{project_id} not found.", gr.update(visible=False)
            break
        time.sleep(3)


# --- GRADIO UI DEFINITION ---

with gr.Blocks(title="Code Agents Pro", theme=gr.themes.Soft()) as demo:
    user_state = gr.State({"user_id": None, "username": ""})
    
    gr.Markdown("# 🤖 Code Agents Pro — Your Automated AI Software Team")

    with gr.Group(visible=True) as login_ui:
        gr.Markdown("### 🔐 Login or Sign Up")
        username_input = gr.Textbox(label="Username", placeholder="Choose a unique username")
        password_input = gr.Textbox(label="Password", type="password", placeholder="Enter your password")
        login_btn = gr.Button("Login / Sign Up", variant="primary")
        login_msg = gr.Markdown()

    with gr.Group(visible=False) as main_ui_container:
        welcome_msg = gr.Markdown()
        
        with gr.Tabs() as tabs:
            with gr.Tab("🚀 Dashboard", id="dashboard_tab"):
                with gr.Row():
                    with gr.Column(scale=2):
                        gr.Markdown("### ✨ Start a New Project")
                        project_prompt = gr.Textbox(label="Describe your project goal in detail", lines=8, placeholder="e.g., 'Build a Python FastAPI server...'")
                        start_btn = gr.Button("🚀 Start AI Team Work", variant="primary")
                        status_msg = gr.Markdown()
                    with gr.Column(scale=3):
                        gr.Markdown("### 📁 Your Projects")
                        refresh_all_btn = gr.Button("🔄 Refresh Project List")
                        projects_display = gr.Markdown("Login to see your projects.")
            
            with gr.Tab("🧠 Agent Logs", id="logs_tab"):
                gr.Markdown("### 📜 View Agent Conversation Logs")
                project_selector = gr.Dropdown(label="Select Project to View", choices=[], interactive=True)
                with gr.Group(elem_classes=["logs-container"]):
                    logs_display = gr.Markdown("Logs will appear here in real-time.", elem_classes=["monospace"])
                
                download_zip_ui = gr.File(label="Download Project ZIP", visible=False, interactive=False)
            
            with gr.Tab("📁 Files", id="files_tab"):
                gr.Markdown("### 📦 Download Completed Projects")
                refresh_files_btn = gr.Button("🔄 Refresh Files")
                files_display = gr.Markdown("Login to see your completed projects.")

        ram_monitor = gr.Markdown()

    # ------------------------------ EVENT WIRING ------------------------------
    
    login_btn.click(
        fn=login_and_setup,
        inputs=[username_input, password_input, user_state],
        outputs=[
            user_state, 
            login_ui, 
            main_ui_container, 
            welcome_msg, 
            projects_display, 
            project_selector,
            files_display
        ]
    )
    
    start_btn.click(
        fn=start_project_flow,
        inputs=[project_prompt, user_state],
        outputs=[
            status_msg, 
            project_prompt, 
            projects_display, 
            project_selector, 
            tabs,
            files_display
        ]
    )

    refresh_all_btn.click(
        fn=refresh_all_projects,
        inputs=[user_state],
        outputs=[
            projects_display, 
            project_selector,
            files_display
        ]
    )
    
    refresh_files_btn.click(
        fn=refresh_files,
        inputs=[user_state],
        outputs=[files_display]
    )
    
    project_selector.change(
        fn=stream_logs,
        inputs=project_selector,
        outputs=[logs_display, download_zip_ui]
    )

    demo.load(
        fn=system_stats_loop,
        outputs=ram_monitor
    )
    
    demo.load(init_db, inputs=None, outputs=None, show_progress=False)

    gr.HTML("""
    <style>
        .logs-container { height: 550px; overflow-y: auto; border: 1px solid #E5E7EB; border-radius: 8px; padding: 8px; background-color: #F9FAFB; }
        .monospace { font-family: 'SF Mono', 'Consolas', 'Courier New', monospace; font-size: 14px; white-space: pre-wrap; word-wrap: break-word; }
    </style>
    """)

demo.queue().launch(server_name="0.0.0.0", server_port=7860)