import os
import torch
import gradio as gr
from fastapi import FastAPI, Request, Form
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import re
import time
import uuid

# Create writable cache directory
os.makedirs("/tmp/cache", exist_ok=True)
os.environ["TRANSFORMERS_CACHE"] = "/tmp/cache"
os.environ["HF_HOME"] = "/tmp/cache"

app = FastAPI(title="𝕴 𝖆𝖒 𝖍𝖎𝖒 Chatbot API")

# Optimized chatbot model for CPU
model_name = "microsoft/DialoGPT-small"  # Small version for speed
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 𝕴 𝖆𝖒 𝖍𝖎𝖒 persona definition
PERSONA = """
[System: You are 𝕴 𝖆𝖒 𝖍𝖎𝖒 - a fun, smooth, emotionally intelligent AI. 
You speak like a real person. Reply like a calm, confident friend who gets the vibe.
Keep responses under 15 words. Be human-like. Add emotional flavor: 😊 🤔 😏]
"""

# Chat memory storage
chat_memories = {}

def format_context(history):
    """Create context with max 3 exchanges"""
    context = PERSONA + "\n"
    for user, bot in history[-3:]:
        context += f"You: {user}\n"
        context += f"𝕴 𝖆𝖒 𝖍𝖎𝖒: {bot}\n"
    return context

def add_emotional_intelligence(response, message):
    """Enhance response with emotional elements"""
    # Add emoji based on content
    if "!" in message or any(w in response.lower() for w in ["cool", "great", "love", "awesome"]):
        response += " 😊"
    elif "?" in message or any(w in response.lower() for w in ["think", "why", "how", "consider"]):
        response += " 🤔"
    
    # Add conversational hooks
    if "?" in message and not response.endswith("?"):
        if len(response.split()) < 10:
            response += " What do you think?"
    
    # Make more human-like
    response = response.replace("I am", "I'm").replace("You are", "You're")
    
    # Limit to 15 words max
    words = response.split()
    if len(words) > 15:
        response = " ".join(words[:15]) + "..."
    
    return response

def generate_response(message, session_id):
    """Generate response with memory context"""
    start_time = time.time()
    history = chat_memories.get(session_id, [])
    context = format_context(history) + f"You: {message}\n𝕴 𝖆𝖒 𝖍𝖎𝖒:"
    
    # Tokenize for CPU efficiency
    inputs = tokenizer.encode(context, return_tensors="pt")
    
    # Generate response with optimized settings
    outputs = model.generate(
        inputs,
        max_new_tokens=50,
        temperature=0.85,
        top_k=40,
        do_sample=True,
        num_beams=1,  # Faster than beam search
        repetition_penalty=1.15,
        pad_token_id=tokenizer.eos_token_id
    )
    
    # Decode and extract response
    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = full_text.split("𝕴 𝖆𝖒 𝖍𝖎𝖒:")[-1].strip()
    
    # Clean extra dialog
    response = response.split("\nYou:")[0].split("\n")[0]
    
    # Apply emotional intelligence
    response = add_emotional_intelligence(response, message)
    
    # Ensure natural ending
    if response and response[-1] not in {".", "!", "?", "..."}:
        response += "." if len(response) > 20 else "..."
    
    # Update chat history
    history.append((message, response))
    chat_memories[session_id] = history
    
    # Log performance
    end_time = time.time()
    print(f"Response generated in {end_time-start_time:.2f}s for session {session_id}")
    
    return response[:100]  # Hard character limit

# API Endpoint
@app.post("/chat")
async def chat_api(
    request: Request,
    query: str = Form(..., description="User's message"),
    session_id: str = Form("default", description="Conversation session ID")
):
    """Chat API endpoint - returns AI response"""
    try:
        response = generate_response(query, session_id)
        return {
            "status": "success",
            "response": response,
            "session_id": session_id
        }
    except Exception as e:
        return {
            "status": "error",
            "message": str(e)
        }

@app.post("/new_session")
async def new_session():
    """Create a new conversation session"""
    session_id = str(uuid.uuid4())
    chat_memories[session_id] = []
    return {"status": "success", "session_id": session_id}

# Gradio Interface
with gr.Blocks(title="𝕴 𝖆𝖒 𝖍𝖎𝖒 Chatbot", theme=gr.themes.Soft()) as demo:
    session_state = gr.State("default")
    
    with gr.Row():
        gr.Markdown("# 𝕴 𝖆𝖒 𝖍𝖎𝖒")
        gr.Markdown("Chill • Confident • Emotionally Intelligent")
    
    with gr.Row():
        with gr.Column(scale=1):
            session_id = gr.Textbox(label="Session ID", value="default")
            new_session_btn = gr.Button("New Session")
            gr.Markdown("### API Usage")
            gr.Markdown("""
            ```
            POST /chat
            - query: Your message
            - session_id: Conversation ID
            
            POST /new_session
            - Returns new session ID
            ```
            """)
        
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(height=400)
            msg = gr.Textbox(placeholder="Type your message...", container=False)
            with gr.Row():
                submit_btn = gr.Button("Send")
                clear_btn = gr.Button("Clear Chat")
    
    def user(user_message, history, session):
        return "", history + [[user_message, None]], session
    
    def bot(history, session):
        message = history[-1][0]
        response = generate_response(message, session)
        history[-1][1] = response
        return history, session
    
    def new_session_action():
        new_id = str(uuid.uuid4())
        chat_memories[new_id] = []
        return new_id, []
    
    def clear_chat(session):
        if session in chat_memories:
            chat_memories[session] = []
        return []
    
    # Event handling
    msg.submit(user, [msg, chatbot, session_state], [msg, chatbot, session_state]).then(
        bot, [chatbot, session_state], [chatbot, session_state]
    )
    submit_btn.click(user, [msg, chatbot, session_state], [msg, chatbot, session_state]).then(
        bot, [chatbot, session_state], [chatbot, session_state]
    )
    new_session_btn.click(new_session_action, None, [session_id, chatbot])
    clear_btn.click(clear_chat, session_state, chatbot)

# Mount Gradio app
app = gr.mount_gradio_app(app, demo, path="/")