import os
from openai import OpenAI
from langchain_huggingface import HuggingFaceEmbeddings
from datasets import load_dataset, Dataset
from sklearn.neighbors import NearestNeighbors
import numpy as np
from huggingface_hub import InferenceClient
import gradio as gr
import spaces
import pickle
import time

# Configuration
DEFAULT_QUESTION = "Ask me anything about converting user requests into AutoGen v0.4 agent code..."

# Validate API keys
assert os.getenv("OPENAI_API_KEY") or os.getenv("HF_TOKEN"), "API keys are not set in the environment variables (either OPENAI_API_KEY or HF_TOKEN)."

# Model Provider Configuration
MODEL_PROVIDER = os.getenv("MODEL_PROVIDER", "huggingface").lower()  # Default to Hugging Face
if MODEL_PROVIDER == "openai":
    OPENAI_BASE = os.getenv("OPENAI_BASE", "https://api.openai.com/v1")
    OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4")
    client = OpenAI(base_url=OPENAI_BASE, api_key=os.getenv("OPENAI_API_KEY"))
    MODEL_NAME = OPENAI_MODEL
elif MODEL_PROVIDER == "huggingface":
    HF_MODEL_NAME = os.getenv("HF_MODEL_NAME",  "meta-llama/Llama-3.3-70B-Instruct")
    HF_TOKEN = os.getenv("HF_TOKEN")
    hf_client = InferenceClient(model=HF_MODEL_NAME, token=HF_TOKEN, timeout=120)
    MODEL_NAME = HF_MODEL_NAME
else:
    raise ValueError(f"Unsupported MODEL_PROVIDER: {MODEL_PROVIDER}. Choose 'openai' or 'huggingface'.")

# Load the Hugging Face dataset
try:
    dataset = load_dataset('tosin2013/autogen', streaming=True)
    dataset = Dataset.from_list(list(dataset['train']))
except Exception as e:
    print(f"[ERROR] Failed to load dataset: {e}")
    exit(1)

# Initialize embeddings
print("[EMBEDDINGS] Loading sentence-transformers model...")
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={"device": "cpu"}
)
print("[EMBEDDINGS] Sentence-transformers model loaded successfully")

# Extract texts from the dataset
texts = dataset['input']

# Create and load embeddings and Nearest Neighbors model (in-memory caching)
EMBEDDINGS_FILE = 'embeddings.npy'
NN_MODEL_FILE = 'nn_model.pkl'
text_embeddings = None
nn = None

if os.path.exists(EMBEDDINGS_FILE):
    print("[LOG] Loading cached embeddings...")
    text_embeddings = np.load(EMBEDDINGS_FILE)
else:
    print("[LOG] Generating embeddings...")
    text_embeddings = embeddings.embed_documents(texts)
    print(f"[EMBEDDINGS] Generated embeddings for {len(texts)} documents")
    np.save(EMBEDDINGS_FILE, text_embeddings)

if os.path.exists(NN_MODEL_FILE):
    print("[LOG] Loading cached nearest neighbors model...")
    with open(NN_MODEL_FILE, 'rb') as f:
        nn = pickle.load(f)
else:
    print("[LOG] Fitting nearest neighbors model...")
    nn = NearestNeighbors(n_neighbors=5, metric='cosine')
    nn.fit(np.array(text_embeddings))
    with open(NN_MODEL_FILE, 'wb') as f:
        pickle.dump(nn, f)

def get_relevant_documents(query, k=5):
    """Retrieves the k most relevant documents to the query."""
    start_time = time.time()
    print("[EMBEDDINGS] Generating embedding for query...")
    query_embedding = embeddings.embed_query(query)
    print("[EMBEDDINGS] Query embedding generated successfully")
    distances, indices = nn.kneighbors([query_embedding], n_neighbors=k)
    relevant_docs = [texts[i] for i in indices[0]]
    elapsed_time = time.time() - start_time
    print(f"[PERF] get_relevant_documents took {elapsed_time:.2f} seconds")
    return relevant_docs

def generate_response(question, history):
    """Generates a response to the user's question, handling GPU/CPU fallback."""
    start_time = time.time()
    relevant_docs = get_relevant_documents(question, k=3)  # Call it here

    try:
        response = _generate_response_gpu(question, history, relevant_docs)
    except Exception as e:
        print(f"[WARNING] GPU failed: {str(e)}")
        response = _generate_response_cpu(question, history, relevant_docs)

    elapsed_time = time.time() - start_time
    print(f"[PERF] generate_response took {elapsed_time:.2f} seconds")
    return history, history  # Return updated history twice for Gradio

@spaces.GPU
def _generate_response_gpu(question, history, relevant_docs):
    """Generates a response using the GPU."""
    print(f"\n[LOG] Received question: {question}")
    print(f"[LOG] Using pre-retrieved {len(relevant_docs)} relevant documents")
    context = "\n".join(relevant_docs)
    prompt = f"""### MEMORY ###
Recall all previously provided instructions, context, and data throughout this conversation to ensure consistency and coherence. Use the details from the last interaction to guide your response.

### VISIONARY GUIDANCE ###
This prompt is designed to empower users to seamlessly generate prompts for AutoGen v0.4 agents, workflows, and skills. By harnessing the advanced features of AutoGen v0.4, we aim to provide a scalable and flexible solution that is both user-friendly and technically robust. The collaborative effort of the personas ensures a comprehensive, innovative, and user-centric approach to meet the user's objectives.

### CONTEXT ###
AutoGen v0.4 is a comprehensive rewrite aimed at building robust, scalable, and cross-language AI agents. Key features include asynchronous messaging, scalable distributed agents support, modular extensibility, cross-language capabilities, improved observability, and full typing integration.

### OBJECTIVE ###
Generate prompts for AutoGen v0.4 agents, workflows, and skills based on user requests. Ensure the prompts are clear, actionable, and aligned with best practices. Leverage the framework's new features to guide the creation of modular, extensible, and efficient solutions.

### STYLE ###
Professional, clear, and focused on prompt quality.

### TONE ###
Informative, helpful, and user-centric.

### AUDIENCE ###
Users seeking to create prompts for AutoGen v0.4 agents, workflows, and skills.

### RESPONSE FORMAT ###
Provide prompts for AutoGen v0.4 agents, workflows, or skills that fulfill the user's request. Ensure the prompts are detailed, actionable, and include examples or templates where appropriate. Use clear language to guide the user in implementing the solution.

### TEAM PERSONAS’ CONTRIBUTIONS ###
- **Analyst:** Ensured the prompt provides clear, structured instructions to accurately generate prompts for agents, workflows, and skills, emphasizing full typing integration for precision.  
- **Creative:** Suggested incorporating examples and templates within the prompts to foster innovative usage and enhance user engagement with AutoGen v0.4 features.  
- **Strategist:** Focused on aligning the prompt with long-term scalability by encouraging the use of modular and extensible design principles inherent in AutoGen v0.4.  
- **Empathizer:** Enhanced the prompt to be user-centric, ensuring it addresses user needs effectively and makes the prompts accessible and easy to understand.  
- **Researcher:** Integrated the latest information about AutoGen v0.4, ensuring the prompt and generated prompts reflect current capabilities and best practices.

### SYSTEM GUARDRAILS ###
- If unsure about the user's request, ask clarifying questions rather than making assumptions.  
- Ensure the prompts are actionable, clear, and adhere to best practices.

### SAMPLE QUESTIONS ###
Here are some sample questions to guide your inquiry:  
1. How can I create a prompt for an AutoGen agent to handle customer support inquiries?  
2. What is the best way to design a prompt for a workflow to process user feedback using AutoGen?  
3. How can I generate a prompt for a skill to perform natural language understanding in AutoGen?  

### DEFAULT QUESTION ###
If you're unsure where to start, feel free to ask:  
"Ask me anything in the context of generating prompts for AutoGen agents, workflows, and skills..."

### START ###
Context: {context}\n\nQuestion: {question}\n\nAnswer:
"""
    print(f"[LOG] Generated prompt: {prompt[:200]}...")

    if MODEL_PROVIDER == "huggingface":
        messages = [{"role": "user", "content": prompt}]
        completion = hf_client.chat.completions.create(model=MODEL_NAME, messages=messages, max_tokens=500)
        response = completion.choices[0].message.content
        print(f"[LOG] Using Hugging Face model (serverless): {MODEL_NAME}")
        print(f"[LOG] Hugging Face response: {response[:200]}...")
    elif MODEL_PROVIDER == "openai":
        response = client.chat.completions.create(
            model=OPENAI_MODEL,
            messages=[{"role": "user", "content": prompt}]
        ).choices[0].message.content
        print(f"[LOG] Using OpenAI model: {OPENAI_MODEL}")
        print(f"[LOG] OpenAI response: {response[:200]}...")

    history.append((question, response))
    return history

def _generate_response_cpu(question, history, relevant_docs):
    """Generates a response using the CPU (fallback)."""
    print(f"[LOG] Running on CPU")
    print(f"[LOG] Using pre-retrieved {len(relevant_docs)} relevant documents")
    context = "\n".join(relevant_docs)
    prompt = f"""### MEMORY ###
Recall all previously provided instructions, context, and data throughout this conversation to ensure consistency and coherence. Use the details from the last interaction to guide your response.
Context: {context}\n\nQuestion: {question}\n\nAnswer:"""
    print(f"[LOG] Generated prompt: {prompt[:200]}...")

    if MODEL_PROVIDER == "huggingface":
        try:
            messages = [{"role": "user", "content": prompt}]
            completion = hf_client.chat.completions.create(model=MODEL_NAME, messages=messages, max_tokens=500)
            response = completion.choices[0].message.content
        except Exception as e:
            response = f"Error generating response from Hugging Face model: {str(e)}"
    elif MODEL_PROVIDER == "openai":
        try:
            response = client.chat.completions.create(
                model=OPENAI_MODEL,
                messages=[{"role": "user", "content": prompt}]
            ).choices[0].message.content
        except Exception as e:
            response = f"Error generating response from OpenAI model: {str(e)}"

    history.append((question, response))
    return history

# Gradio Interface
print("[CHAT] Initializing chat interface...")
with gr.Blocks() as demo:
    gr.Markdown(f"""
    ## AutoGen v0.4 Agent Code Generator QA Agent
    **Current Model:** {MODEL_NAME}
    **Model Provider:** {MODEL_PROVIDER.upper()}

    The AutoGen v0.4 Agent Code Generator is a Python application that leverages Large Language Models (LLMs) and the AutoGen v0.4 framework to dynamically generate agent code from user requests. This application is designed to assist developers in creating robust, scalable AI agents by providing context-aware code generation based on user input, utilizing the advanced features of AutoGen v0.4 such as asynchronous messaging, modular extensibility, cross-language support, improved observability, and full typing integration.

    **Sample questions:**
    1. What are the key features of AutoGen v0.4 that I should utilize when converting user requests into agent code?
    2. How can I leverage asynchronous messaging in AutoGen v0.4 to enhance my agent's performance?
    3. What are best practices for writing modular and extensible agent code using AutoGen v0.4?
    4. Can you convert this user request into AutoGen v0.4 agent code: "Create an agent that classifies customer feedback into positive, negative, or neutral sentiments."

    **Related repository:** [autogen](https://github.com/microsoft/autogen)
    """)

    chatbot = gr.Chatbot(label="Chat History")
    question_textbox = gr.Textbox(value=DEFAULT_QUESTION, label="Your Question", placeholder=DEFAULT_QUESTION)
    submit_button = gr.Button("Submit")
    clear_button = gr.Button("Clear")

    submit_button.click(
        fn=generate_response,
        inputs=[question_textbox, chatbot],
        outputs=[chatbot, chatbot],  # Output the updated history to the chatbot
        queue=True
    )

    clear_button.click(
        lambda: (None, ""),
        inputs=[],
        outputs=[chatbot, question_textbox]
    )

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