import os
from typing import List, Dict

import openai
from qdrant_client import QdrantClient
from qdrant_client.http import models as qdrant_models
from dotenv import load_dotenv

load_dotenv()

# Configuration
QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY", None)
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

if not OPENAI_API_KEY:
    raise RuntimeError("OPENAI_API_KEY environment variable not set")

openai.api_key = OPENAI_API_KEY

# Collection name
COLLECTION_NAME = "faq_collection"

# Embedding dimension for text-embedding-ada-002
EMBEDDING_DIM = 1536

# Sample FAQ data
FAQ_DATA = [
    {
        "question": "What is the return policy?",
        "answer": "You can return any item within 30 days of purchase."
    },
    {
        "question": "How do I track my order?",
        "answer": "Use the tracking link sent to your email after shipping."
    },
    {
        "question": "Do you offer international shipping?",
        "answer": "Yes, we ship to most countries worldwide."
    },
    {
        "question": "What payment methods are accepted?",
        "answer": "We accept credit cards, PayPal, and bank transfers."
    },
    {
        "question": "How can I contact customer support?",
        "answer": "Email us at support@example.com or call 1-800-123-4567."
    }
]

def get_embedding(text: str) -> List[float]:
    """
    Generate an embedding for the given text using OpenAI's embedding model.
    """
    response = openai.Embedding.create(
        input=text,
        model="text-embedding-ada-002"
    )
    return response["data"][0]["embedding"]

def create_or_recreate_collection(client: QdrantClient) -> None:
    """
    Create a new collection or recreate it if it already exists.
    """
    client.recreate_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=qdrant_models.VectorParams(
            size=EMBEDDING_DIM,
            distance=qdrant_models.Distance.COSINE
        )
    )

def ingest_faqs(client: QdrantClient, faqs: List[Dict[str, str]]) -> None:
    """
    Ingest FAQ data into Qdrant.
    """
    points = []
    for idx, faq in enumerate(faqs):
        vector = get_embedding(faq["question"])
        point = qdrant_models.PointStruct(
            id=idx,
            vector=vector,
            payload={
                "question": faq["question"],
                "answer": faq["answer"]
            }
        )
        points.append(point)

    # Upsert points in batches
    batch_size = 100
    for i in range(0, len(points), batch_size):
        batch = points[i:i+batch_size]
        client.upsert(
            collection_name=COLLECTION_NAME,
            points=batch
        )

def query_faq(client: QdrantClient, question: str, top_k: int = 1) -> str:
    """
    Query the FAQ collection for the most relevant answer.
    """
    query_vector = get_embedding(question)
    search_result = client.search(
        collection_name=COLLECTION_NAME,
        query_vector=query_vector,
        limit=top_k,
        with_payload=True
    )
    if not search_result:
        return "Sorry, I couldn't find an answer to your question."
    # Return the answer from the top result
    return search_result[0].payload.get("answer", "Answer not found.")

def main() -> None:
    client = QdrantClient(
        host=QDRANT_HOST,
        port=QDRANT_PORT,
        api_key=QDRANT_API_KEY
    )

    # Ingest FAQs (only if collection is empty or you want to refresh)
    print("Ingesting FAQ data into Qdrant...")
    create_or_recreate_collection(client)
    ingest_faqs(client, FAQ_DATA)
    print("Ingestion complete.")

    print("\nFAQ Bot is ready. Type your question (or 'exit' to quit).")
    while True:
        user_input = input("\nYour question: ").strip()
        if user_input.lower() in {"exit", "quit"}:
            print("Goodbye!")
            break
        answer = query_faq(client, user_input)
        print(f"Answer: {answer}")

if __name__ == "__main__":
    main()