import os
import logging
import numpy as np
from typing import List, Optional
from langchain_community.vectorstores import FAISS
from langchain.embeddings.base import Embeddings
from gradio_client import Client
import gradio as gr

# Configuration
DATA_FILE = "data-mtc.txt"  # This file is no longer used in the Space
DATABASE_DIR = "."  # Database files are in the root directory
CHUNK_SIZE = 800
TOP_K_RESULTS = 100
SIMILARITY_THRESHOLD = 0.4

BASE_SYSTEM_PROMPT = """
Répondez en français selon ces règles :

1. Utilisez EXCLUSIVEMENT le contexte fourni
2. Structurez la réponse en :
   - Définition principale
   - Caractéristiques clés (3 points maximum)
   - Relations avec d'autres concepts
3. Si aucune information pertinente, indiquez-le clairement

Contexte :
{context}
"""

# Logging configuration
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler()  # Output to console in the Space
    ]
)


class GradioEmbeddings(Embeddings):
    """Embedding management using Gradio API"""

    def __init__(self):
        super().__init__()
        self.client = Client("localsavageai/embijiji3")

    def _generate_embedding(self, text: str) -> np.ndarray:
        """Generate an embedding via the Gradio API"""
        try:
            result = self.client.predict(
                document=text.strip(),
                api_name="/embed"
            )
            if not isinstance(result, list):
                raise ValueError("Invalid embedding response from Gradio API")
            return np.array(result, dtype=np.float32)
        except Exception as e:
            logging.error(f"Embedding error: {str(e)}")
            raise RuntimeError("Failed to generate embedding") from e

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return [self._generate_embedding(text).tolist() for text in texts]

    def embed_query(self, text: str) -> List[float]:
        return self._generate_embedding(text).tolist()


def initialize_vector_store() -> FAISS:
    """Robust initialization of the vector store"""
    embeddings = GradioEmbeddings()

    try:
        logging.info("Loading existing database...")
        return FAISS.load_local(
            DATABASE_DIR,
            embeddings,
            allow_dangerous_deserialization=True
        )
    except Exception as e:
        logging.error(f"FAISS loading error: {str(e)}")
        raise


def generate_response(user_input: str, vector_store: FAISS) -> Optional[str]:
    """Generate a response with complete error handling"""
    try:
        docs_scores = vector_store.similarity_search_with_score(
            user_input,
            k=TOP_K_RESULTS * 3
        )

        filtered_docs = [
            (doc, score) for doc, score in docs_scores
            if score < SIMILARITY_THRESHOLD
        ]
        filtered_docs.sort(key=lambda x: x[1])

        if not filtered_docs:
            return ("No matches found in MTC texts. "
                    "Try using more specific terms.")

        best_docs = [doc for doc, _ in filtered_docs[:TOP_K_RESULTS]]

        context = "\n".join(
            f"=== Source {i + 1} ===\n{doc.page_content}\n"
            for i, doc in enumerate(best_docs)
        )

        response = Client("Qwen/Qwen2.5-Max-Demo").predict(
            query=user_input,
            history=[],
            system=BASE_SYSTEM_PROMPT.format(context=context),
            api_name="/model_chat"
        )

        if isinstance(response, tuple) and len(response) >= 2:
            chat_history = response[1]
            if isinstance(chat_history, list) and len(chat_history) > 0:
                last_message = chat_history[-1]
                if isinstance(last_message, (list, tuple)) and len(last_message) >= 2:
                    return last_message[1]

        return "Response unavailable - Please rephrase your question."

    except Exception as e:
        logging.error(f"Generation error: {str(e)}", exc_info=True)
        return "An error occurred while generating the response."


def chatbot(query):
    """Main function to run the chatbot"""
    try:
        vs = initialize_vector_store()
        response = generate_response(query, vs)
        return response or "No response generated."
    except Exception as e:
        logging.error(f"Chatbot error: {str(e)}")
        return f"An error occurred: {str(e)}"


# Gradio Interface
if __name__ == "__main__":
    try:
        interface = gr.Interface(
            fn=chatbot,
            inputs=gr.Textbox(lines=7, placeholder="Enter your query here..."),
            outputs=gr.Textbox(lines=7, placeholder="Response from MTC will appear here..."),
            title="MTC Chatbot",
            description="Ask questions about MTC and get answers based on the provided data."
        )
        interface.launch()

    except Exception as e:
        logging.critical(f"CRITICAL ERROR: {str(e)}")
        print("Failed to launch Gradio interface.  Check logs.")