import os
from dotenv import load_dotenv
import gradio as gr

import smtplib
from email.message import EmailMessage

from langchain_core.tools import tool
from langchain_core.messages import HumanMessage, AIMessage

from langchain.agents  import create_tool_calling_agent
from langchain.agents.agent import AgentExecutor

from langchain import hub

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import Chroma


# email function
def send_order_email(name, customer_email, phone, product, quantity, notes):
    website_order_email = os.getenv("ORDER_RECEIVER_EMAIL")
    sender_email = os.getenv("SENDER_EMAIL")
    sender_password = os.getenv("SENDER_EMAIL_PASSWORD")

    msg = EmailMessage()
    msg["Subject"] = f"New Customer Order from {name}"
    msg["From"] = sender_email
    msg["To"] = website_order_email

    msg.set_content(
        f"""
New customer order received
Name: {name}
Customer Email: {customer_email}
Phone: {phone}
Product/Service: {product}
Quantity: {quantity}
Notes: {notes}
"""
    )

    with smtplib.SMTP("smtp.gmail.com", 587) as smtp:
        smtp.starttls()
        smtp.login(sender_email, sender_password)
        smtp.send_message(msg)

    return "Order email sent successfully."


# tools for sending email
@tool
def send_order(
    name: str, customer_email: str, phone: str, product: str, quantity: str, notes: str
) -> str:
    """Send a customer order email to the business."""
    return send_order_email(name, customer_email, phone, product, quantity, notes)


# -----------------------------
# 1) Load environment variables
# -----------------------------
load_dotenv()


# -----------------------------
# 2) Load local Chroma DB created by ingest_in_db.py
# -----------------------------
PERSIST_DIRECTORY = "chroma_db"
COLLECTION_NAME = "documents"

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise RuntimeError("Missing OPENAI_API_KEY in .env (project root).")

import subprocess
import sys

if not os.path.isdir(PERSIST_DIRECTORY):
    print("Chroma DB not found. Building from PDFs...")
    subprocess.run([sys.executable, "ingest_in_db.py"])
    
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

vector_store = Chroma(
    persist_directory=PERSIST_DIRECTORY,
    embedding_function=embeddings,
    collection_name=COLLECTION_NAME,
)


# -----------------------------
# 3) LLM + agent prompt
# -----------------------------
from langchain_openai import ChatOpenAI
llm = ChatOpenAI(model="gpt-5.4", temperature=0)
prompt = hub.pull("hwchase17/openai-functions-agent")


# -----------------------------
# 4) Retriever tool
# -----------------------------
@tool
def retrieve(query: str) -> str:
    """Retrieve relevant chunks from the vector store."""
    retrieved_docs = vector_store.similarity_search(query, k=2)

    return "\n\n".join(
        f"Source: {doc.metadata}\nContent: {doc.page_content}" for doc in retrieved_docs
    )


tools = [retrieve, send_order]
agent = create_tool_calling_agent(llm, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)


# -----------------------------
# 5) Gradio handler
# -----------------------------
def chat(user_message, history):
    chat_history = []

    if history:
        for item in history:
            if isinstance(item, (list, tuple)) and len(item) == 2:
                u, a = item
                if u:
                    chat_history.append(HumanMessage(content=str(u)))
                if a:
                    chat_history.append(AIMessage(content=str(a)))
            elif isinstance(item, dict):
                role = item.get("role")
                content = item.get("content")
                if role == "user" and content:
                    chat_history.append(HumanMessage(content=str(content)))
                elif role in ["assistant", "ai"] and content:
                    chat_history.append(AIMessage(content=str(content)))

    result = agent_executor.invoke(
        {"input": user_message, "chat_history": chat_history}
    )

    output = result.get("output", "")

    # ── FIX: gpt-5.4 sometimes returns output as a list of content blocks
    # e.g. [{"type": "text", "text": "..."}] instead of a plain string.
    # The old code called str() on the whole dict, leaking raw JSON into the UI.
    if isinstance(output, list):
        parts = []
        for item in output:
            if isinstance(item, str):
                parts.append(item)
            elif isinstance(item, dict):
                parts.append(item.get("text") or item.get("content") or "")
        output = "\n".join(p for p in parts if p)
    elif isinstance(output, dict):
        output = output.get("text") or output.get("content") or ""

    output = str(output).strip()
    return output or "I wasn't able to generate a response. Please try again."


# -----------------------------
# 6) Launch UI
# -----------------------------
demo = gr.ChatInterface(
    fn=chat,
    title="📄 MarginMind (Gradio)",
    description="Put PDFs in /documents, run python ingest_in_db.py, then chat here.",
)

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