import gradio as gr
import torch
import sys
import subprocess

# Install required packages
def install_packages():
    packages = ["peft==0.7.0", "sentencepiece==0.1.99", "protobuf==3.20.3"]
    for package in packages:
        try:
            __import__(package.split('==')[0])
            print(f"✅ {package} already installed")
        except ImportError:
            print(f"📦 Installing {package}...")
            subprocess.check_call([sys.executable, "-m", "pip", "install", package])

install_packages()

from transformers import T5Tokenizer, T5ForConditionalGeneration
from peft import PeftModel

# Configuration
PEFT_MODEL_ID = "daffaaditya/daffa-ai"
BASE_MODEL = "t5-small"

print("=" * 50)
print(f"🚀 Loading PEFT Adapter: {PEFT_MODEL_ID}")
print(f"📦 Base Model: {BASE_MODEL}")
print("=" * 50)

try:
    # Load tokenizer dari BASE MODEL, bukan dari adapter
    print("1. Loading tokenizer from base model...")
    tokenizer = T5Tokenizer.from_pretrained(BASE_MODEL)
    
    # Load base model
    print("2. Loading base model...")
    base_model = T5ForConditionalGeneration.from_pretrained(BASE_MODEL)
    
    # Load adapter
    print("3. Loading adapter...")
    model = PeftModel.from_pretrained(base_model, PEFT_MODEL_ID)
    
    # Merge adapter (optional, tapi lebih cepat untuk inference)
    print("4. Merging adapter with base model...")
    model = model.merge_and_unload()
    
    print("✅ SUCCESS: PEFT model loaded and merged!")
    
except Exception as e:
    print(f"❌ Error: {e}")
    print("🔄 Falling back to plain T5-small without adapter...")
    tokenizer = T5Tokenizer.from_pretrained(BASE_MODEL)
    model = T5ForConditionalGeneration.from_pretrained(BASE_MODEL)

# Setup device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"📱 Device: {device}")
model.to(device)
model.eval()

def generate_code(instruction, max_length=150, temperature=0.7):
    """Generate Python code from instruction"""
    try:
        # Clean instruction
        instruction = instruction.strip()
        
        # Format prompt lebih baik
        if "prima" in instruction.lower():
            prompt = "Write a Python function to check if a number is prime:"
        elif "faktorial" in instruction.lower():
            prompt = "Write a Python function to calculate factorial:"
        elif "reverse" in instruction.lower():
            prompt = "Write a Python function to reverse a string:"
        elif "lingkaran" in instruction.lower() or "circle" in instruction.lower():
            prompt = "Write a Python function to calculate circle area:"
        else:
            prompt = f"Write Python code for: {instruction}"
        
        # Tokenize
        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=128).to(device)
        
        # Generate
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_length=max_length,
                temperature=temperature,
                do_sample=True,
                num_return_sequences=1,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
                repetition_penalty=1.3,
                no_repeat_ngram_size=3
            )
        
        # Decode
        generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Clean output
        if generated.startswith(prompt):
            generated = generated[len(prompt):].strip()
        
        # Jika output kosong atau pendek
        if len(generated) < 10:
            # Fallback templates
            fallbacks = {
                "prima": """def is_prime(n):
    if n < 2:
        return False
    for i in range(2, int(n**0.5) + 1):
        if n % i == 0:
            return False
    return True""",
                "faktorial": """def factorial(n):
    if n == 0:
        return 1
    return n * factorial(n-1)""",
                "reverse": """def reverse_string(s):
    return s[::-1]"""
            }
            
            for key, code in fallbacks.items():
                if key in instruction.lower():
                    return code
        
        return generated
        
    except Exception as e:
        return f"# Error\n# {str(e)}\n\nPlease try a different prompt."

# Create Gradio Interface
with gr.Blocks(title="Daffa AI Coder") as demo:
    gr.Markdown("# 🐍 Daffa AI Coder")
    gr.Markdown("Fine-tuned T5 model for Python code generation")
    
    with gr.Row():
        with gr.Column():
            instruction = gr.Textbox(
                label="Instruction (Indonesian/English)",
                placeholder="Example: buat fungsi untuk cek bilangan prima",
                lines=3
            )
            
            with gr.Accordion("⚙️ Settings", open=False):
                max_length = gr.Slider(50, 300, value=150, label="Max Length")
                temperature = gr.Slider(0.1, 1.5, value=0.7, label="Temperature")
            
            generate_btn = gr.Button("Generate Code", variant="primary")
        
        with gr.Column():
            output = gr.Code(
                label="Generated Python Code",
                language="python",
                lines=15
            )
    
    # Examples
    examples = [
        ["buat fungsi untuk cek bilangan prima"],
        ["create a function to calculate factorial"],
        ["write a function to reverse a string"],
        ["fungsi untuk menghitung luas lingkaran"],
        ["function to convert celsius to fahrenheit"]
    ]
    
    gr.Examples(
        examples=examples,
        inputs=instruction,
        outputs=output,
        fn=lambda x: generate_code(x, 150, 0.7),
        label="Try these examples:"
    )
    
    # Footer
    gr.Markdown("---")
    gr.Markdown(f"**Model:** {PEFT_MODEL_ID} | **Base:** {BASE_MODEL} | **Status:** {'Adapter Loaded' if 'PeftModel' in str(type(model)) else 'Base Model Only'}")
    
    # Events
    generate_btn.click(generate_code, [instruction, max_length, temperature], output)

demo.launch()