app.py

import os
import gradio as gr
from huggingface_hub import hf_hub_download
import importlib.util

# --- Private Repository Information ---
PRIVATE_DATASET_ID = os.getenv("PRIVATE_DATASET_ID")
HF_TOKEN = os.getenv("HF_TOKEN")
INDEX_SUBDIR = os.getenv("INDEX_SUBDIR", ".")

# --- Core Logic Download and Import ---
try:
    AGENT_CODE_PATH = hf_hub_download(
        repo_id=PRIVATE_DATASET_ID,
        filename="deepv_core.py",
        repo_type="dataset",
        token=HF_TOKEN
    )
    spec = importlib.util.spec_from_file_location("deepv_core_module", AGENT_CODE_PATH)
    agent_module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(agent_module)
    run_generation = agent_module.run_generation
    
except Exception as e:
    def show_error(*args):
        return f"// ERROR: Failed to load core agent code. Details: {e}", "", ""

# --- Gradio UI setup below ---
with gr.Blocks(title="DeepV for RTL (Model-Agnostic)", theme=gr.themes.Soft()) as demo:
    gr.Markdown("## DeepV for RTL Code Generation — Model-Agnostic (Bring Your Own API Key)", elem_id="main-title")

    with gr.Row():
        with gr.Column(scale=2):
            with gr.Row():
                model_choice = gr.Dropdown(
                    choices=["gpt-4o", "gpt-4o-mini", "gpt-4.1", "gpt-5-chat-latest"],
                    value="gpt-4o",
                    label="Model"
                )
                api_key = gr.Textbox(label="OpenAI API Key", type="password", placeholder="sk-...")

            gr.Markdown(
                """
                **Note:** Your API key is used for the current session only and is not saved or stored.
                """
            )

            spec = gr.Textbox(
                label="Design Specification (natural language or I/O contract)",
                placeholder="e.g., 8-bit UART transmitter with baud rate generator ...",
                lines=10,
                elem_id="spec-input"
            )
            with gr.Row():
                use_rag = gr.Checkbox(value=True, label="Use Retrieval (RAG)")
                top_k = gr.Slider(1, 5, value=3, step=1, label="Top-K retrieved examples")

            # --- New UI with three text boxes for generation settings ---
            with gr.Row():
                temperature_tb = gr.Textbox(label="Temperature", value="0.2", scale=1)
                top_p_tb = gr.Textbox(label="Top-p", value="0.9", scale=1)
                max_new_tokens_tb = gr.Textbox(label="Max tokens", value="768", scale=1)

            # --- Loading State Components ---
            with gr.Row():
                run_btn = gr.Button("Generate Verilog", variant="primary", elem_id="generate-button")
                loading_state = gr.Markdown(
                    value="Generating...",
                    visible=False,
                    elem_id="loading-state"
                )
            
        with gr.Column(scale=3):
            gr.Markdown("**Output**")
            with gr.Group():
                out_code = gr.Textbox(
                    label="Generated Verilog",
                    lines=28,
                    interactive=False,
                    placeholder="// Your Verilog code will appear here",
                    elem_id="verilog-output"
                )
                copy_button = gr.Button("📋", variant="secondary", elem_id="copy-button")
            
            with gr.Row():
                clear_btn = gr.ClearButton(
                    value="Clear All",
                    components=[spec, out_code]
                )

    # --- Wrapper function to handle output from the generation agent ---
    def generate_only(
        spec, use_rag, top_k, model_choice, api_key, temperature, top_p, max_new_tokens
    ):
        verilog_code, _, _ = run_generation(
            spec, use_rag, top_k, model_choice, api_key, temperature, top_p, max_new_tokens
        )
        return verilog_code
    
    def copy_to_clipboard_fn(text):
        return text

    def show_loading():
        return [gr.update(visible=False), gr.update(visible=True)]

    def hide_loading():
        return [gr.update(visible=True), gr.update(visible=False)]

    run_btn.click(
        fn=show_loading,
        inputs=[],
        outputs=[run_btn, loading_state],
        show_progress=False
    ).then(
        fn=generate_only,
        inputs=[spec, use_rag, top_k, model_choice, api_key, temperature_tb, top_p_tb, max_new_tokens_tb],
        outputs=[out_code],
    ).then(
        fn=hide_loading,
        inputs=[],
        outputs=[run_btn, loading_state]
    )

    clear_btn.click(fn=lambda: "Ready", outputs=[])
    spec.submit(fn=lambda: "Ready", outputs=[])

    copy_button.click(
        fn=copy_to_clipboard_fn,
        inputs=[out_code],
        outputs=[],
        js="""
            (text) => {
                const el = document.createElement('textarea');
                el.value = text;
                document.body.appendChild(el);
                el.select();
                document.execCommand('copy');
                document.body.removeChild(el);
            }
        """
    )

    demo.css = """
        @import url('https://fonts.googleapis.com/css2?family=Poppins:wght@600&display=swap');

        #main-title {
            font-family: 'Poppins', sans-serif;
            font-weight: 600;
            letter-spacing: 1px;
        }

        #verilog-output > label > .label-wrap {
            position: relative;
        }

        #copy-button {
            position: absolute;
            top: 20px;
            right: 10px;
            z-index: 1000;
            background-color: #F0F0F0;
            color: #333;
            border-radius: 5px;
            border: none;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        }

        #loading-state {
            padding: 14px;
            text-align: center;
            font-size: 1.2em;
            font-weight: 800; 
            color: #1E3A8A;
            background-color: #e6f2ff;
            border-radius: 8px;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
            animation: pulse 1s infinite;
        }

        @keyframes pulse {
            0% { background-color: #e6f2ff; }
            50% { background-color: #d2e4f7; }
            100% { background-color: #e6f2ff; }
        }
    """

if __name__ == "__main__":
    if 'agent_module' in locals():
        demo.launch()
    else:
        with gr.Blocks() as error_demo:
            gr.Markdown("# Initialization Error")
            gr.Markdown(f"An error occurred while loading the application code. Please check your configuration.")
            gr.Textbox(label="Error Details", value=str(e), lines=5)
        error_demo.launch()