import gradio as gr
import os
import traceback
import spaces  # <--- ही नवीन ओळ (ZeroGPU साठी)
from vajra.assembler import assemble, AssembleError
from vajra.linker import link, LinkerError
from vajra.cpu import VajraCPU
from vajra.utils import format_registers, format_log, load_example_program


# <--- हा डेकोरेटर (महत्त्वाचा!)
@spaces.GPU
def run_vajra(source: str, mode: str, max_steps: int):
    """Execute Vajra program and return results."""
    try:
        # 1. Assemble
        imem, data, directives, symbols = assemble(source)
        
        # 2. Link (load HF datasets if needed)
        data, symbols = link(directives, data, symbols)
        
        # 3. Create CPU and load program
        cpu = VajraCPU()
        cpu.load_program(imem)
        cpu.load_data(data)
        
        # 4. Execute
        if mode == "Run":
            result = cpu.run(max_steps=max_steps)
            output = f"✅ Execution complete.\n\nFinal R0: {result:.6f}\nCycles: {cpu.cycle_count}\n\n"
            output += format_registers(cpu.R, cpu.PC, cpu.SP)
            output += f"\n\n--- Execution Log ---\n{format_log(cpu.log)}"
            
        elif mode == "Step":
            steps = min(max_steps, 50)
            for _ in range(steps):
                if cpu.halted:
                    break
                cpu.step()
            
            output = f"⏯ Stepped {cpu.cycle_count} instructions.\n\n"
            output += format_registers(cpu.R, cpu.PC, cpu.SP)
            output += f"\n\n--- Execution Log ---\n{format_log(cpu.log)}"
            
        else:  # Profile
            cpu.run(max_steps=max_steps)
            output = f"📊 Profile:\n"
            output += f"  Total cycles: {cpu.cycle_count}\n"
            output += f"  Halted: {cpu.halted}\n"
            output += f"  Final PC: {cpu.PC}\n"
            output += f"  Final R0: {cpu.R[0]:.6f}\n"
        
        return output
        
    except (AssembleError, LinkerError) as e:
        return f"❌ Error: {e}"
    except Exception as e:
        return f"❌ Runtime error: {e}\n\nDetails:\n{traceback.format_exc()}"


def load_example(example_name: str):
    """Load example program."""
    return load_example_program(example_name)


# Gradio Interface
with gr.Blocks(title="Vajra-ISA Virtual Machine") as demo:
    gr.Markdown("# 🕉️ Vajra-ISA (वज्र-ISA)")
    gr.Markdown("Vedic-inspired Virtual Machine for Neuro-Symbolic AI")
    
    with gr.Row():
        with gr.Column(scale=2):
            code_input = gr.Code(
                label="Vajra Program (.vajra)",
                language="python",
                value=load_example_program('addition'),
                lines=20
            )
            
            with gr.Row():
                example_dropdown = gr.Dropdown(
                    choices=['addition', 'multiply', 'perceptron', 'dhyana', 'memory'],
                    value='addition',
                    label="Load Example"
                )
                mode_dropdown = gr.Dropdown(
                    choices=['Run', 'Step', 'Profile'],
                    value='Run',
                    label="Execution Mode"
                )
            
            max_steps = gr.Slider(
                minimum=10,
                maximum=100000,
                value=1000,
                step=10,
                label="Max Steps"
            )
            
            run_btn = gr.Button("▶ Execute", variant="primary")
        
        with gr.Column(scale=2):
            output = gr.Textbox(
                label="Output & Register State",
                lines=30,
                interactive=False
            )
    
    # Event handlers
    example_dropdown.change(
        fn=load_example,
        inputs=example_dropdown,
        outputs=code_input
    )
    
    run_btn.click(
        fn=run_vajra,
        inputs=[code_input, mode_dropdown, max_steps],
        outputs=output
    )

if __name__ == "__main__":
    # Gradio 6.0 compatible launch
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        css="footer {display: none !important;}",
        debug=True
    )