import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

def load_model():
    repo_id = "forestav/gguf_lora_model"
    model_file = "unsloth.F16.gguf"  
    
    local_path = hf_hub_download(
        repo_id=repo_id,
        filename=model_file
    )
    
    print(f"Loading model from: {local_path}")
    
    model = Llama(
        model_path=local_path,
        n_ctx=2048,
        n_threads=8
    )
    
    return model

def generate_instructions(input_text, instruction_type, complexity, audience):
    # Craft a comprehensive system prompt
    system_prompt = f"""You are an expert at creating clear, precise instructions. 
    Generate instructions that are:
    - Type: {instruction_type}
    - Complexity Level: {complexity}
    - Target Audience: {audience}
    
    Core Input Context: {input_text}
    
    Guidelines:
    - Use clear, step-by-step language
    - Ensure instructions are actionable and specific
    - Include safety warnings or prerequisites if relevant
    - Adapt complexity to the specified audience level"""

    # Prepare messages for instruction generation
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"Please generate comprehensive instructions for: {input_text}"}
    ]
    
    # Generate response
    response = model.create_chat_completion(
        messages=messages,
        max_tokens=1024,
        temperature=0.7,
        top_p=0.95,
    )
    
    return response['choices'][0]['message']['content']

# Load model globally
print("Starting model loading...")
model = load_model()
print("Model loaded successfully!")

# Create Gradio interface
demo = gr.Blocks(title="Instruction Craft AI")

with demo:
    gr.Markdown("# 📝 Instruction Crafting Assistant")
    gr.Markdown("Generate precise, tailored instructions for any task or process.")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Describe the task or process")
            
            instruction_type = gr.Dropdown(
                label="Instruction Type",
                choices=[
                    "How-to Guide", 
                    "Technical Manual", 
                    "Safety Procedure", 
                    "Educational Tutorial", 
                    "Cooking Recipe", 
                    "DIY Project", 
                    "Professional Workflow"
                ]
            )
            
            complexity = gr.Dropdown(
                label="Complexity Level",
                choices=[
                    "Beginner", 
                    "Intermediate", 
                    "Advanced", 
                    "Expert"
                ]
            )
            
            audience = gr.Dropdown(
                label="Target Audience",
                choices=[
                    "Children", 
                    "Students", 
                    "General Public", 
                    "Professionals", 
                    "Experts"
                ]
            )
            
            generate_btn = gr.Button("Craft Instructions", variant="primary")
        
        with gr.Column():
            output_text = gr.Textbox(label="Generated Instructions", lines=20)
    
    generate_btn.click(
        fn=generate_instructions, 
        inputs=[input_text, instruction_type, complexity, audience], 
        outputs=output_text
    )

    # Set default example values
    input_text.value = "Change a car tire"
    instruction_type.value = "How-to Guide"
    complexity.value = "Intermediate"
    audience.value = "General Public"

# Launch the interface
demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    share=False
)