import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import re

# ... (il codice di caricamento del modello rimane lo stesso) ...
model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

# --- VERSIONE FINALE DELLA LOGICA DI GENERAZIONE E PARSING ---
def generate_code_and_files(user_prompt):
    # Continuiamo a usare il prompt "few-shot" perché produce il codice migliore
    messages = [
        {
            "role": "system",
            "content": "You are a silent code generation machine. You MUST ONLY output code. Do not write any explanation. Your entire response must be ONLY the code files requested by the user, enclosed in the specified format. The format is: <file name=\"FILE_NAME\">file content</file>"
        },
        { "role": "user", "content": "Create a simple hello world app in Python Flask." },
        { "role": "assistant", "content": "<file name=\"app.py\">\nfrom flask import Flask\n\napp = Flask(__name__)\n\n@app.route('/')\ndef hello():\n    return 'Hello, World!'\n\nif __name__ == '__main__':\n    app.run(debug=True)\n</file>\n<file name=\"requirements.txt\">\nflask\n</file>" },
        { "role": "user", "content": user_prompt }
    ]
    
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    print("Inizio la generazione del codice...")
    inputs = tokenizer(prompt, return_tensors="pt")
    output = model.generate(
        **inputs,
        max_new_tokens=1024,
        do_sample=True,
        temperature=0.1,
        top_p=0.9
    )
    generated_text = tokenizer.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    print("Testo grezzo generato:\n", generated_text)

    # --- PARSER FINALE E AGGRESSIVO ---
    files = []
    
    # Tentativo 1: Il formato ideale <file>
    ideal_matches = re.findall(r'<file name="([^"]+)">(.*?)</file>', generated_text, re.DOTALL)
    if ideal_matches:
        print("Formato ideale <file> trovato!")
        files = [{"name": name, "content": content} for name, content in ideal_matches]

    # Tentativo 2: Se non trova il formato ideale, usiamo la forza bruta
    if not files:
        print("Formato ideale non trovato. Applico il post-processing aggressivo.")
        
        # Estrai tutto il codice Python dentro i blocchi ```python ... ```
        python_code_match = re.search(r'```python\n(.*?)\n```', generated_text, re.DOTALL)
        if python_code_match:
            app_py_content = python_code_match.group(1).strip()
            files.append({"name": "app.py", "content": app_py_content})
            
            # Cerca le dipendenze nel codice (es. 'from flask import')
            # Questo è un parser molto semplice ma spesso efficace
            requirements_content = ""
            if "from flask import" in app_py_content:
                requirements_content += "flask\n"
            if "import pandas" in app_py_content:
                requirements_content += "pandas\n"
            if "import numpy" in app_py_content:
                requirements_content += "numpy\n"
            
            if requirements_content:
                 files.append({"name": "requirements.txt", "content": requirements_content.strip()})

    if not files:
        return f"Il modello non ha generato i file in un formato riconoscibile. Ecco l'output grezzo:\n\n```\n{generated_text}\n```"

    # Formattiamo l'output per Gradio
    formatted_output = f"✅ Ho generato {len(files)} file (usando il post-processing):\n\n"
    for file_info in files:
        name = file_info["name"]
        content = file_info["content"]
        formatted_output += f"### 📄 `{name}`\n"
        formatted_output += f"```\n{content}\n```\n\n"

    return formatted_output

# --- INTERFACCIA GRADIO (rimane uguale) ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🤖 App Code Generator AI (CPU Edition)")
    gr.Markdown("Descrivi l'applicazione che vuoi creare. Questa versione usa un modello più piccolo (TinyLlama) e gira su CPU.")
    
    with gr.Row():
        user_input = gr.Textbox(
            label="Descrivi la tua applicazione",
            lines=5,
            placeholder="Es: 'Crea una app Flask che mostra data e ora attuali in formato HH:MM.'"
        )
    
    generate_btn = gr.Button("🚀 Genera Codice", variant="primary")
    output_display = gr.Markdown(label="File Generati")
    
    generate_btn.click(
        fn=generate_code_and_files, 
        inputs=user_input, 
        outputs=output_display
    )

demo.launch()