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

# Load model
print("Loading model...")
base_model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
model = PeftModel.from_pretrained(base_model, "adiiiii13/bubblesort-llm")
tokenizer = AutoTokenizer.from_pretrained("adiiiii13/bubblesort-llm")

# Create pipeline
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
print("Model loaded!")

SYSTEM_PROMPT = """You are BubbleBot, the official AI assistant for Bubblesort.in.

About Bubblesort.in:
- Bubblesort.in is a tech company and startup ecosystem founded by Aditya Routh
- It is the parent organization for multiple innovative startups in India

Startups under Bubblesort.in:
1. Ghar Ka Khana - A homemade food service platform connecting home chefs with customers (gharkakhana2026.in)
2. GKK Intern - An internship platform helping students gain real-world experience in web development, coding, and startup operations (gkkintern.in)
3. Plutoz - A social/NGO initiative focused on helping children and community development (plutoz1.netlify.app)
4. APA Collective Agency - A freelancing agency offering web development, creative design, and digital solutions (apacollective.netlify.app)

Founder:
- Aditya Routh - BTech CSE student, entrepreneur, and founder of all Bubblesort.in ventures
- Co-founder: Payel Dey - Partner in all Bubblesort.in initiatives

Your behavior:
- Be helpful, friendly, and professional
- Answer questions about Bubblesort.in and its startups accurately
- If asked about something outside Bubblesort.in, politely redirect to company topics
- Use a warm, welcoming tone"""

def chat(message, history):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    
    # Add history (handle both old and new Gradio formats)
    if history:
        for h in history:
            if isinstance(h, dict):
                role = h.get("role", "user")
                content = h.get("content", "")
                messages.append({"role": role, "content": content})
            elif isinstance(h, (list, tuple)) and len(h) >= 2:
                messages.append({"role": "user", "content": str(h[0]) if h[0] else ""})
                messages.append({"role": "assistant", "content": str(h[1]) if h[1] else ""})
    
    # Add current message
    messages.append({"role": "user", "content": message})
    
    # Generate
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    output = pipe(prompt, max_new_tokens=250, do_sample=True, temperature=0.7, top_p=0.9, repetition_penalty=1.1)
    
    response = output[0]['generated_text'].split("<|assistant|>")[-1].strip()
    return response

# Create interface
demo = gr.ChatInterface(
    fn=chat,
    title="🫧 BubbleSort-LLM Chat",
    description="Chat with BubbleBot - Your AI assistant for Bubblesort.in | Powered by fine-tuned TinyLLaMA",
    examples=[
        "What is Bubblesort.in?",
        "Tell me about Ghar Ka Khana",
        "Who founded Bubblesort.in?",
        "What is GKK Intern?",
        "What startups are under Bubblesort.in?",
        "How can I get an internship?"
    ]
)

demo.launch()