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

# --- 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}", "", ""

# --- Base64 Encoding for Logo ---
def get_logo_base64():
    try:
        with open("DeepV_logo.png", "rb") as image_file:
            return base64.b64encode(image_file.read()).decode("utf-8")
    except FileNotFoundError:
        print("Logo file not found. Ensure DeepV_logo.png is in the correct directory.")
        return ""

logo_base64 = get_logo_base64()

# --- Gradio UI setup below ---
with gr.Blocks(title="DeepV for RTL (Model-Agnostic)", theme=gr.themes.Soft()) as demo:
    # The logo is now at the top middle, outside of any column structure
    if logo_base64:
        gr.HTML(f"""
            <div id="deepv-logo-container">
                <img id="deepv-logo" src="data:image/png;base64,{logo_base64}" alt="DeepV Logo">
            </div>
        """)
    else:
        gr.Markdown("Logo not found.")

    with gr.Row():
        with gr.Column(scale=2):
            with gr.Row():
                model_choice = gr.Dropdown(
                    choices=["gpt-4o", "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 RAG")
                top_k = gr.Slider(1, 5, value=3, step=1, label="Top-K retrieved examples")

            # --- Generation settings as text boxes ---
            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):
            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_no_anim():
        return [gr.update(visible=False), gr.update(visible=True)]

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

    run_btn.click(
        fn=show_loading_no_anim,
        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_no_anim,
        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');

        /* New CSS to make the layout more compact */
        .gradio-container {
            margin-top: 0 !important;
        }

        #deepv-logo-container {
            display: flex;
            justify-content: center;
            align-items: center;
            padding: 0;
            margin-bottom: 0px; 
            margin-top: -35px; /* Pull the logo up */
        }

        #deepv-logo {
            width: 450px; 
            height: 125px; 
            object-fit: contain;
        }
        
        #verilog-output {
            position: relative;
        }

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

        #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()