# app.py

import os
import pandas as pd
import chardet
import logging
import gradio as gr
from typing import Optional, List, Tuple, ClassVar, Dict

from sentence_transformers import SentenceTransformer, util, CrossEncoder
from langchain.llms.base import LLM
import google.generativeai as genai

# Import smolagents components
from smolagents import CodeAgent, LiteLLMModel, DuckDuckGoSearchTool, ManagedAgent

###############################################################################
# 1) Logging Setup
###############################################################################
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("Daily Wellness AI")

###############################################################################
# 2) API Key Handling and Enhanced GeminiLLM Class
###############################################################################
def clean_api_key(key: str) -> str:
    """Remove non-ASCII characters and strip whitespace from the API key."""
    return ''.join(c for c in key if ord(c) < 128).strip()

# Load the GEMINI API key from environment variables
gemini_api_key = os.environ.get("GEMINI_API_KEY")

if not gemini_api_key:
    logger.error("GEMINI_API_KEY environment variable not set.")
    raise EnvironmentError("Please set the GEMINI_API_KEY environment variable.")

gemini_api_key = clean_api_key(gemini_api_key)
logger.info("GEMINI API Key loaded successfully.")

# Configure Google Generative AI
try:
    genai.configure(api_key=gemini_api_key)
    logger.info("Configured Google Generative AI with provided API key.")
except Exception as e:
    logger.error(f"Failed to configure Google Generative AI: {e}")
    raise e

class GeminiLLM(LLM):
    model_name: ClassVar[str] = "gemini-2.0-flash-exp"
    temperature: float = 0.7
    top_p: float = 0.95
    top_k: int = 40
    max_tokens: int = 2048

    @property
    def _llm_type(self) -> str:
        return "custom_gemini"

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        generation_config = {
            "temperature": self.temperature,
            "top_p": self.top_p,
            "top_k": self.top_k,
            "max_output_tokens": self.max_tokens,
        }

        try:
            logger.debug(f"Initializing GenerativeModel with config: {generation_config}")
            model = genai.GenerativeModel(
                model_name=self.model_name,
                generation_config=generation_config,
            )
            logger.debug("GenerativeModel initialized successfully.")

            chat_session = model.start_chat(history=[])
            logger.debug("Chat session started.")

            response = chat_session.send_message(prompt)
            logger.debug(f"Prompt sent to model: {prompt}")
            logger.debug(f"Raw response received: {response.text}")

            return response.text
        except Exception as e:
            logger.error(f"Error generating response with GeminiLLM: {e}")
            logger.debug("Exception details:", exc_info=True)
            raise e

# Instantiate the GeminiLLM globally
llm = GeminiLLM()

###############################################################################
# 3) CSV Loading and Processing
###############################################################################
def load_csv(file_path: str):
    try:
        if not os.path.isfile(file_path):
            logger.error(f"CSV file does not exist: {file_path}")
            return [], []
        
        with open(file_path, 'rb') as f:
            result = chardet.detect(f.read())
            encoding = result['encoding']

        data = pd.read_csv(file_path, encoding=encoding)
        if 'Question' not in data.columns or 'Answers' not in data.columns:
            raise ValueError("CSV must contain 'Question' and 'Answers' columns.")
        data = data.dropna(subset=['Question', 'Answers'])

        logger.info(f"Loaded {len(data)} entries from {file_path}")
        return data['Question'].tolist(), data['Answers'].tolist()
    except Exception as e:
        logger.error(f"Error loading CSV: {e}")
        return [], []

# Path to your CSV file (ensure 'AIChatbot.csv' is in the repository)
csv_file_path = "AIChatbot.csv"
corpus_questions, corpus_answers = load_csv(csv_file_path)
if not corpus_questions:
    raise ValueError("Failed to load the knowledge base.")

###############################################################################
# 4) Sentence Embeddings & Cross-Encoder
###############################################################################
embedding_model_name = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
try:
    embedding_model = SentenceTransformer(embedding_model_name)
    logger.info(f"Loaded embedding model: {embedding_model_name}")
except Exception as e:
    logger.error(f"Failed to load embedding model: {e}")
    raise e

try:
    question_embeddings = embedding_model.encode(corpus_questions, convert_to_tensor=True)
    logger.info("Encoded question embeddings successfully.")
except Exception as e:
    logger.error(f"Failed to encode question embeddings: {e}")
    raise e

cross_encoder_name = "cross-encoder/ms-marco-MiniLM-L-6-v2"
try:
    cross_encoder = CrossEncoder(cross_encoder_name)
    logger.info(f"Loaded cross-encoder model: {cross_encoder_name}")
except Exception as e:
    logger.error(f"Failed to load cross-encoder model: {e}")
    raise e

###############################################################################
# 5) Retrieval + Re-Ranking
###############################################################################
class EmbeddingRetriever:
    def __init__(self, questions, answers, embeddings, model, cross_encoder):
        self.questions = questions
        self.answers = answers
        self.embeddings = embeddings
        self.model = model
        self.cross_encoder = cross_encoder

    def retrieve(self, query: str, top_k: int = 3):
        try:
            query_embedding = self.model.encode(query, convert_to_tensor=True)
            scores = util.pytorch_cos_sim(query_embedding, self.embeddings)[0].cpu().tolist()
            scored_data = sorted(zip(self.questions, self.answers, scores), key=lambda x: x[2], reverse=True)[:top_k]

            cross_inputs = [[query, candidate[0]] for candidate in scored_data]
            cross_scores = self.cross_encoder.predict(cross_inputs)

            reranked = sorted(zip(scored_data, cross_scores), key=lambda x: x[1], reverse=True)
            final_retrieved = [(entry[0][1], entry[1]) for entry in reranked]
            logger.debug(f"Retrieved and reranked answers: {final_retrieved}")
            return final_retrieved
        except Exception as e:
            logger.error(f"Error during retrieval: {e}")
            logger.debug("Exception details:", exc_info=True)
            return []

retriever = EmbeddingRetriever(corpus_questions, corpus_answers, question_embeddings, embedding_model, cross_encoder)

###############################################################################
# 6) Sanity Check Tool
###############################################################################
class QuestionSanityChecker:
    def __init__(self, llm: GeminiLLM):
        self.llm = llm

    def is_relevant(self, question: str) -> bool:
        prompt = (
            f"You are an assistant that determines whether a question is relevant to daily wellness.\n\n"
            f"Question: {question}\n\n"
            f"Is the above question relevant to daily wellness? Respond with 'Yes' or 'No' only."
        )
        try:
            response = self.llm._call(prompt)
            is_yes = 'yes' in response.lower()
            logger.debug(f"Sanity check response: {response}, interpreted as {is_yes}")
            return is_yes
        except Exception as e:
            logger.error(f"Error in sanity check: {e}")
            logger.debug("Exception details:", exc_info=True)
            return False

# Instantiate the sanity checker globally
sanity_checker = QuestionSanityChecker(llm)

###############################################################################
# 7) smolagents Integration: GROQ Model and Web Search
###############################################################################
# Initialize the smolagents' LiteLLMModel with GROQ model
smol_model = LiteLLMModel("groq/llama3-8b-8192")

# Instantiate the DuckDuckGo search tool
search_tool = DuckDuckGoSearchTool()

# Create the web agent with the search tool
web_agent = CodeAgent(
    tools=[search_tool],
    model=smol_model
)

# Define the managed web agent
managed_web_agent = ManagedAgent(
    agent=web_agent,
    name="web_search",
    description="Runs a web search for you. Provide your query as an argument."
)

# Create the manager agent with managed web agent and additional tools if needed
manager_agent = CodeAgent(
    tools=[],  # Add additional tools here if required
    model=smol_model,
    managed_agents=[managed_web_agent]
)

###############################################################################
# 8) Answer Expansion
###############################################################################
class AnswerExpander:
    def __init__(self, llm: GeminiLLM):
        self.llm = llm

    def expand(self, query: str, retrieved_answers: List[str]) -> str:
        try:
            reference_block = "\n".join(f"- {idx+1}) {ans}" for idx, ans in enumerate(retrieved_answers, start=1))
            prompt = (
                f"You are Daily Wellness AI, a friendly wellness expert. Below are multiple "
                f"potential answers retrieved from a local knowledge base. You have a user question.\n\n"
                f"Question: {query}\n\n"
                f"Retrieved Answers:\n{reference_block}\n\n"
                "Please synthesize these references into a single cohesive, creative, and brand-aligned response. "
                "Add practical tips and positivity, and end with a short inspirational note. "
                "Please provide a concise response in no more than 4 sentences.\n\n"
                "Disclaimer: This is general wellness information, not a substitute for professional medical advice."
            )
            logger.debug(f"Generated prompt for answer expansion: {prompt}")
            response = self.llm._call(prompt)
            logger.debug(f"Expanded answer: {response}")
            return response.strip()
        except Exception as e:
            logger.error(f"Error expanding answer: {e}")
            logger.debug("Exception details:", exc_info=True)
            return "Sorry, an error occurred while generating a response."

answer_expander = AnswerExpander(llm)

###############################################################################
# 9) Query Handling
###############################################################################
def handle_query(query: str) -> str:
    if not query or not isinstance(query, str) or len(query.strip()) == 0:
        return "Please provide a valid question."

    try:
        # Sanity Check: Determine if the question is relevant to daily wellness
        is_relevant = sanity_checker.is_relevant(query)
        if not is_relevant:
            return "Your question seems out of context or not related to daily wellness. Please ask a wellness-related question."

        # Proceed with retrieval
        retrieved = retriever.retrieve(query)
        if not retrieved:
            return "I'm sorry, I couldn't find an answer to your question."

        # Check similarity threshold
        top_score = retrieved[0][1]  # Assuming the list is sorted descending
        similarity_threshold = 0.3  # Adjust this threshold based on empirical results

        if top_score < similarity_threshold:
            # Perform web search using manager_agent
            logger.info("Similarity score below threshold. Performing web search.")
            web_search_response = manager_agent.run(query)
            logger.debug(f"Web search response: {web_search_response}")

            # Optionally, process the web_search_response if needed
            # For simplicity, return the web search response directly
            return (
                f"**Daily Wellness AI**\n\n"
                f"{web_search_response}\n\n"
                "Disclaimer: This information is retrieved from the web and is not a substitute for professional medical advice.\n\n"
                "Wishing you a calm and wonderful day!"
            )

        # Proceed with answer expansion using retrieved_answers
        responses = [ans[0] for ans in retrieved]
        expanded_answer = answer_expander.expand(query, responses)
        return expanded_answer
    except Exception as e:
        logger.error(f"Error handling query: {e}")
        logger.debug("Exception details:", exc_info=True)
        return "An error occurred while processing your request."

###############################################################################
# 10) Gradio Interface
###############################################################################
def gradio_interface(query: str):
    try:
        response = handle_query(query)
        formatted_response = response  # Response is already formatted
        return formatted_response
    except Exception as e:
        logger.error(f"Error in Gradio interface: {e}")
        logger.debug("Exception details:", exc_info=True)
        return "**An error occurred while processing your request. Please try again later.**"

interface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Textbox(
        lines=2,
        placeholder="e.g., What is box breathing?",
        label="Ask Daily Wellness AI"
    ),
    outputs=gr.Markdown(label="Answer from Daily Wellness AI"),
    title="Daily Wellness AI",
    description="Ask wellness-related questions and receive synthesized, creative answers.",
    theme="default",
    examples=[
        "What is box breathing and how does it help reduce anxiety?",
        "Provide a daily wellness schedule incorporating box breathing techniques.",
        "What are some tips for maintaining good posture while working at a desk?",
        "Who is the CEO of Hugging Face?"  # Example of an out-of-context question
    ],
    allow_flagging="never"
)

###############################################################################
# 11) Launch Gradio
###############################################################################
if __name__ == "__main__":
    try:
        # For Hugging Face Spaces, set share=False
        interface.launch(server_name="0.0.0.0", server_port=7860, debug=False)
    except Exception as e:
        logger.error(f"Failed to launch Gradio interface: {e}")
        logger.debug("Exception details:", exc_info=True)