import os
import requests
from openai import OpenAI

# Configuration
FLTR_API_KEY = os.getenv("FLTR_API_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
FLTR_BASE_URL = "https://api.fltr.com/v1"

fltr_headers = {
    "Authorization": f"Bearer {FLTR_API_KEY}",
    "Content-Type": "application/json"
}

openai_client = OpenAI(api_key=OPENAI_API_KEY)

# Step 1: Create a dataset
def create_dataset(name: str, description: str):
    response = requests.post(
        f"{FLTR_BASE_URL}/datasets",
        headers=fltr_headers,
        json={
            "name": name,
            "description": description,
            "is_public": False
        }
    )
    response.raise_for_status()
    return response.json()["id"]

# Step 2: Upload documents
def upload_document(dataset_id: str, content: str, metadata: dict):
    response = requests.post(
        f"{FLTR_BASE_URL}/datasets/{dataset_id}/documents",
        headers=fltr_headers,
        json={
            "content": content,
            "metadata": metadata
        }
    )
    response.raise_for_status()
    return response.json()

# Step 3: Query FLTR for relevant context
def search_knowledge_base(dataset_id: str, query: str, limit: int = 3):
    response = requests.post(
        f"{FLTR_BASE_URL}/mcp/query",
        headers=fltr_headers,
        json={
            "query": query,
            "dataset_id": dataset_id,
            "limit": limit
        }
    )
    response.raise_for_status()
    return response.json()["results"]

# Step 4: Generate answer with GPT-4
def generate_answer(question: str, context_chunks: list):
    # Format context from retrieved chunks
    context = "\n\n".join([
        f"[Source: {chunk['metadata'].get('title', 'Unknown')}]\n{chunk['content']}"
        for chunk in context_chunks
    ])

    # Create prompt with context
    messages = [
        {
            "role": "system",
            "content": "You are a helpful assistant. Answer questions based on the provided context. If the context doesn't contain enough information, say so."
        },
        {
            "role": "user",
            "content": f"Context:\n{context}\n\nQuestion: {question}"
        }
    ]

    # Call GPT-4
    response = openai_client.chat.completions.create(
        model="gpt-4",
        messages=messages,
        temperature=0.7,
        max_tokens=500
    )

    return response.choices[0].message.content

# Step 5: Main RAG function
def answer_question(dataset_id: str, question: str):
    print(f"Question: {question}\n")

    # Retrieve relevant context
    print("Searching knowledge base...")
    chunks = search_knowledge_base(dataset_id, question, limit=3)

    print(f"Found {len(chunks)} relevant chunks\n")

    # Generate answer
    print("Generating answer...")
    answer = generate_answer(question, chunks)

    print(f"\nAnswer: {answer}\n")

    # Show sources
    print("Sources:")
    for i, chunk in enumerate(chunks, 1):
        title = chunk['metadata'].get('title', 'Unknown')
        score = chunk['score']
        print(f"{i}. {title} (relevance: {score:.2f})")

    return answer

# Example usage
if __name__ == "__main__":
    # Create dataset
    dataset_id = create_dataset(
        name="Product Documentation",
        description="FLTR product docs for RAG demo"
    )

    print(f"Created dataset: {dataset_id}\n")

    # Upload sample documents
    docs = [
        {
            "content": "FLTR supports three authentication methods: API keys for services, OAuth 2.1 for MCP clients, and session tokens for web apps. API keys provide 1,000 requests per hour.",
            "metadata": {"title": "Authentication Guide", "category": "security"}
        },
        {
            "content": "FLTR uses hybrid search combining vector embeddings with keyword matching. You can enable Cohere reranking for even better results. The default embedding model is text-embedding-3-small.",
            "metadata": {"title": "Search Guide", "category": "features"}
        },
        {
            "content": "To integrate FLTR with Zapier, use the Webhooks by Zapier action. Set the URL to https://api.fltr.com/v1/mcp/query and include your API key in the Authorization header.",
            "metadata": {"title": "Zapier Integration", "category": "integrations"}
        }
    ]

    for doc in docs:
        upload_document(dataset_id, doc["content"], doc["metadata"])
        print(f"Uploaded: {doc['metadata']['title']}")

    print("\nWaiting for indexing to complete...\n")
    import time
    time.sleep(3)  # Give FLTR time to process

    # Ask questions
    questions = [
        "How do I authenticate with FLTR?",
        "What search methods does FLTR support?",
        "How can I use FLTR with Zapier?"
    ]

    for question in questions:
        print("=" * 60)
        answer_question(dataset_id, question)
        print()