from langchain.tools import Tool
from langchain_community.retrievers import BM25Retriever
from langchain.docstore.document import Document
from langchain_core.messages import HumanMessage
import datasets
from langchain_openai import AzureChatOpenAI
import os
from dotenv import load_dotenv

load_dotenv()

# Create LLM instance once
conversation_llm = AzureChatOpenAI(
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
    api_key=os.getenv("AZURE_OPENAI_API_KEY"),
    deployment_name=os.getenv("DEPLOYMENT_NAME"),
    openai_api_version=os.getenv("OPENAI_API_VERSION"),
    temperature=0.75,
    streaming=False,
    verbose=False
)


def load_guest_dataset():
    # Load the dataset
    guest_dataset = datasets.load_dataset(
        "agents-course/unit3-invitees", split="train")

    # Convert dataset entries into Document objects
    docs = [
        Document(
            page_content="\n".join([
                f"Name: {guest['name']}",
                f"Relation: {guest['relation']}",
                f"Description: {guest['description']}",
                f"Email: {guest['email']}"
            ]),
            metadata={"name": guest["name"]}
        )
        for guest in guest_dataset
    ]

    return docs


docs = load_guest_dataset()
bm25_retriever = BM25Retriever.from_documents(docs)


def generate_conversation_starter(description: str) -> str:
    """Generate a conversation starter based on guest description"""
    try:
        generate_prompt = (
            f"Generate a very simple and short conversation starter from the description of the person.\n\n"
            f"For example:\n"
            f"Description: Rear Admiral Grace Hopper was a trailblazer in computer programming and helped invent the first compiler. "
            f"She's passionate about teaching and loves telling stories about debugging.\n\n"
            f"Conversation Starter: Ask her about the time she found a real bug in a computer — she loves that story!\n\n"
            f"Description: {description}\n\n"
            f"Conversation Starter:"
        )

        response = conversation_llm.invoke(
            [HumanMessage(content=generate_prompt)])
        return response.content.strip()
    except Exception:
        return "Ask them about their background and interests!"


def retrieve_info_from_name(query: str) -> str:
    """Retrieves detailed information about gala guests based on their name or relation."""
    results = bm25_retriever.invoke(query)
    if results:
        guest_info_with_starters = []

        for i, doc in enumerate(results[:3], 1):
            guest_info = doc.page_content

            # Extract description from the content
            lines = guest_info.split('\n')
            description = ""
            for line in lines:
                if line.startswith("Description:"):
                    description = line.replace("Description:", "").strip()
                    break

            # Add guest info
            result_text = f"Guest {i}:\n{guest_info}"

            # Add conversation starter if description exists
            if description:
                conversation_starter = generate_conversation_starter(
                    description)
                result_text += f"\n💬 Conversation Starter: {conversation_starter}"

            guest_info_with_starters.append(result_text)

        return "\n\n" + "="*50 + "\n\n".join(guest_info_with_starters)
    else:
        return "No matching guest information found."


guest_info_tool = Tool(
    name="guest_info_retriever",
    func=retrieve_info_from_name,
    description="Retrieves detailed information about gala guests based on their name or relation, including conversation starters."
)