import os, datetime, math, numexpr, logging, requests, uuid
import gradio as gr

from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_core.documents import Document
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_retrieval_chain
from langchain_core.tools import StructuredTool
from langchain_core.messages import SystemMessage, HumanMessage
from langchain.prompts import PromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate, SystemMessagePromptTemplate
from langchain_community.vectorstores import FAISS

# === SETUP ===
GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")

# === LLM ===
llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash")

# === TOOLS ===
def dani_drinks() -> str:
    """Dani's favorite drinks info on beers, wine and water."""
    return (
    "Dani's favourite drinks:\n"
    "Beers: Peroni (better if 33cl), Moretti, Ichnusa (unfiltered)\n"
    "Wine: Red fruity, White cold only, No Rosè\n"
    "Water: Normal or sparkling, only cold\n"
    )

def calculator(expression: str) -> str:
    """Math calculations."""
    local_dict = {"pi": math.pi, "e": math.e}
    result = str(numexpr.evaluate(expression.strip(), local_dict=local_dict))
    return result

# Nuovo Weather Tool
def weather_tool(location: str) -> str:
    """Return current weather for the given city name using Open‑Meteo APIs."""
    # 1. Geocoding: city → lat, lon
    geocode = requests.get(
        "https://geocoding-api.open-meteo.com/v1/search",
        params={"name": location, "count": 1}
    )
    geo_json = geocode.json()
    if not geo_json.get("results"):
        return f"Non ho trovato la città '{location}'. Riprova."

    first = geo_json["results"][0]
    lat = first["latitude"]
    lon = first["longitude"]
    name = first.get("name")
    country = first.get("country", "")
    
    # 2. Current weather
    weather_resp = requests.get(
        "https://api.open-meteo.com/v1/forecast",
        params={
            "latitude": lat,
            "longitude": lon,
            "current_weather": True,
            "timezone": "auto"
        }
    )
    w = weather_resp.json().get("current_weather")
    if not w:
        return f"Meteo non disponibile per {location}."

    temp = w.get("temperature")
    wind = w.get("windspeed")
    return (f"Meteo attuale a {name}, {country}: {temp}°C, vento {wind} km/h")

# === RETRIEVAL TOOL VIA FAISS ===
# create mock docs and embed
docs = [
    Document(page_content="Milan was founded in 1900."),
    Document(page_content="Inter was founded in 1950."),
    Document(page_content="Juve was founded in 1920."),
    Document(page_content="Torino was founded in 1910."),
]
embeddings = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004")
vector_db = FAISS.from_documents(docs, embeddings)
retriever = vector_db.as_retriever()

rag_prompt = PromptTemplate.from_template(
    "Answer the questions using ONLY the documents.\nQuestion: {input}\n\nDocuments:\n{context}"
)

doc_chain = create_stuff_documents_chain(llm, rag_prompt)
rag_chain = create_retrieval_chain(retriever, doc_chain)

def retrieval_tool(query: str) -> str:
    """Answer questions about italian soccer teams"""
    result = rag_chain.invoke(
        {"input": query},
    )
    return result["answer"]
    
tools = [
    StructuredTool.from_function(dani_drinks),
    StructuredTool.from_function(calculator),
    StructuredTool.from_function(weather_tool),
    StructuredTool.from_function(retrieval_tool)
]

# === AGENT ===
system_prompt = SystemMessagePromptTemplate.from_template(
    "You are DaniBot, an agent that can answer math questions, tell info about Dani's favorite drinks and tell the weather in a city."
    "You can also answer questions about italian soccer teams "
    "Dani is your boss. Use the tools if needed. Answer with an Italian accent! "
    "You can also answer general base knowledge questions if you know the answer"
    "If the user seems wanting to quit, thank him/her and say 'SUCA!'"
)

human_prompt = HumanMessagePromptTemplate.from_template("{input}\n{agent_scratchpad}")
chat_prompt = ChatPromptTemplate.from_messages([system_prompt, human_prompt])

# === AGENT ===
raw_agent = create_tool_calling_agent(llm, tools, chat_prompt)

agent_executor = AgentExecutor(
    agent=raw_agent,
    tools=tools,
    verbose=False
)

# === GRADIO INTERFACCIA + window ===
def chat_fn(message, history):
    session_id = str(uuid.uuid4())
    result = agent_executor.invoke(
        {"input": message},
        config={"configurable": {"session_id": session_id}}
    )
    return result["output"]

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # DaniBot 🤖

        Questo chatbot può rispondere sui drinks di Dani, ma anche a domande di matematica, informazioni sulle squadre di calcio italiane,
        il meteo, e altro ancora. Usa dei tools integrati per darti risposte più accurate.

        Scrivi un messaggio per iniziare la chat!
        Scrivi 'q' o 'quit' per terminare la chat!
        """
    )
    chatbot = gr.ChatInterface(
    	fn=chat_fn,
    	examples=["Quando è stata fondata la Juventus?"],
    	title="DaniBot - ask me something!",
    	type="messages"
    )

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