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

base_model_id = "NousResearch/Llama-2-7b-chat-hf"

lora_path = "Arsh014/lora-llama2-finetuned"



tokenizer = AutoTokenizer.from_pretrained(base_model_id)

model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    load_in_8bit=True,
    torch_dtype=torch.float16,
    device_map="auto"
)

try:
    model = PeftModel.from_pretrained(model, lora_path)
    model.eval() # Set model to evaluation mode
except Exception as e:
    print(f"Error loading LoRA adapter from {lora_path}. Ensure it exists and is correct.")
    print(f"Error: {e}")

# 5. Create a text-generation pipeline
print("Creating text-generation pipeline.")
pipe = pipeline(
    "text-generation", 
    model=model, 
    tokenizer=tokenizer
)

def format_prompt(instruction, code):
    """Formats the instruction and input code into the required prompt template."""
    return f"""### Instruction:
{instruction}

### Input:
{code}

### Response:"""

def explain_dockerfile(instruction, code):
    """Generates the explanation using the text-generation pipeline."""
    if not instruction or not code:
        return "Please provide both an instruction and the Dockerfile code."

    prompt = format_prompt(instruction, code)

    # Generate response
    response = pipe(
        prompt, 
        max_new_tokens=256, 
        do_sample=True, 
        temperature=0.7,
        return_full_text=False 
    )
    

    generated_text = response[0]["generated_text"].strip()
    
    if "### Response:" in generated_text:
        return generated_text.split("### Response:")[-1].strip()
    
    return generated_text

# 6. Gradio Interface
print("Launching Gradio Interface...")
iface = gr.Interface(
    fn=explain_dockerfile,
    inputs=[
        gr.Textbox(
            label="Instruction", 
            placeholder="e.g., Explain the function of each line and the overall goal of this Dockerfile.",
            value="Explain this Dockerfile in detail and suggest one security improvement.",
            lines=2
        ),
        gr.Textbox(
            label="Dockerfile Code", 
            lines=10, 
            placeholder="Paste your Dockerfile here, e.g., \nFROM python:3.9-slim\nWORKDIR /app\nCOPY requirements.txt .\nRUN pip install -r requirements.txt\nCOPY . .\nCMD [\"python\", \"app.py\"]",
            value="FROM node:18-alpine\nWORKDIR /usr/src/app\nCOPY package*.json ./ \nRUN npm install\nCOPY . .\nEXPOSE 3000\nCMD [ \"npm\", \"start\" ]"
        )
    ],
    outputs=gr.Textbox(
        label="Explanation (Generated by LoRA Model)", 
        lines=15
    ),
    title="LoRA-Tuned Llama-2 Dockerfile Explainer",
    description="A simple application to explain complex Dockerfiles using a fine-tuned Llama-2 model (via LoRA).",
    live=False
)

if __name__ == "__main__":
    iface.launch()