import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel
import torch

print("🔄 Model yükleniyor...")

# Model yükle
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16
)

base_model = AutoModelForCausalLM.from_pretrained(
    "bigcode/starcoder2-3b",
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)

model = PeftModel.from_pretrained(base_model, "MuratKomurcu/starcoder2-stm32-turkish")
tokenizer = AutoTokenizer.from_pretrained("MuratKomurcu/starcoder2-stm32-turkish")

print("✅ Model hazır!")

def generate_code(project_type, pin_config, parameters, temperature):
    """STM32 kodu üret"""
    
    templates = {
        "LED Blink": "Write complete STM32F401RE LED blink code with main function",
        "UART Communication": "Write complete STM32F401RE UART communication code with main function",
        "ADC Reading": "Write complete STM32F401RE ADC reading code with main function",
        "PWM Motor Control": "Write complete STM32F401RE PWM motor control code with main function",
        "I2C Sensor": "Write complete STM32F401RE I2C sensor code with main function",
        "SPI Communication": "Write complete STM32F401RE SPI communication code with main function",
        "Timer Interrupt": "Write complete STM32F401RE timer interrupt code with main function",
        "Temperature Sensor": "Write complete STM32F401RE temperature sensor reading code with main function",
        "Servo Control": "Write complete STM32F401RE servo control code with main function",
        "Watchdog Timer": "Write complete STM32F401RE watchdog timer code with main function",
    }
    
    instruction = templates.get(project_type, project_type)
    input_text = f"{pin_config}, {parameters}, include SystemClock_Config"
    
    prompt = f"""### Instruction:
{instruction}

### Input:
{input_text}

### Response:
"""
    
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(
        **inputs,
        max_new_tokens=2048,
        temperature=temperature,
        top_p=0.95,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id
    )
    
    code = tokenizer.decode(outputs[0], skip_special_tokens=True)
    code = code.split("### Response:")[-1].strip()
    
    return code

# GRADIO UI
with gr.Blocks(title="STM32 Code Generator", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    # 🔧 STM32 HAL Code Generator
    ### AI-Powered Embedded Code Generation for STM32F4
    
    This model generates complete STM32F4 HAL code with **85% accuracy** on realistic projects.
    Fine-tuned on 55,000 STM32 examples.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Configuration")
            
            project = gr.Dropdown(
                choices=[
                    "LED Blink",
                    "UART Communication", 
                    "ADC Reading",
                    "PWM Motor Control",
                    "I2C Sensor",
                    "SPI Communication",
                    "Timer Interrupt",
                    "Temperature Sensor",
                    "Servo Control",
                    "Watchdog Timer"
                ],
                label="Project Type",
                value="LED Blink"
            )
            
            pins = gr.Textbox(
                label="Pin Configuration",
                placeholder="Example: LED on PA5, Button on PC13",
                value="LED on PA5"
            )
            
            params = gr.Textbox(
                label="Additional Parameters",
                placeholder="Example: blink every 1000ms, with debounce",
                value="blink every 1000ms"
            )
            
            temp = gr.Slider(
                0.1, 1.0,
                value=0.3,
                step=0.1,
                label="Temperature (creativity)",
                info="Lower = deterministic, Higher = creative"
            )
            
            btn = gr.Button("🚀 Generate Code", variant="primary", size="lg")
        
        with gr.Column(scale=2):
            output = gr.Code(
                label="Generated STM32 Code",
                language="c",
                lines=35
            )
    
    gr.Markdown("### 💡 Example Projects")
    gr.Examples(
        examples=[
            ["LED Blink", "LED on PA5", "blink every 1000ms", 0.1],
            ["UART Communication", "USART2 on PA2 PA3", "115200 baud, transmit Hello World every second", 0.3],
            ["Temperature Sensor", "ADC1 Channel 0 on PA0", "LM35 sensor (10mV per °C), display via USART2 every 1000ms", 0.1],
            ["PWM Motor Control", "PWM TIM2 Channel 1 on PA5, potentiometer ADC1 on PA0", "control motor speed 0-100%, PWM frequency 1kHz", 0.1],
            ["I2C Sensor", "I2C1", "read from device 0x48, transmit via USART2 115200 baud every 1000ms", 0.1],
        ],
        inputs=[project, pins, params, temp]
    )
    
    btn.click(
        fn=generate_code,
        inputs=[project, pins, params, temp],
        outputs=output
    )
    
    gr.Markdown("""
    ---
    **Model:** [MuratKomurcu/starcoder2-stm32-turkish](https://huggingface.co/MuratKomurcu/starcoder2-stm32-turkish)  
    **Performance:** 17/20 tests passed (85% accuracy)  
    **Base:** bigcode/starcoder2-3b fine-tuned on 55K STM32 examples  
    **Created by:** Murat Kömürcü
    """)

demo.launch()