# services/general_agent.py

import logging
import os
import faiss
from fastapi import HTTPException
from langchain_ollama import OllamaLLM
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from sentence_transformers import SentenceTransformer

from services.cache_service import CacheService


_embedder = None
def get_embedder():
    global _embedder
    if _embedder is None:
        _embedder = SentenceTransformer("/opt/models/bge-m3", device="cpu")
    return _embedder
embedder = get_embedder()

vec_path = "vector_store/faiss.index"
cache_service = CacheService()


def _load_index() -> faiss.IndexFlatL2:
    if os.path.exists(vec_path):
        return faiss.read_index(vec_path)
    dim = embedder.get_sentence_embedding_dimension()
    return faiss.IndexFlatL2(dim)


logger = logging.getLogger(__name__)


class GeneralAgent:
    """Handles general conversational queries in multiple languages."""

    def __init__(self, llm: OllamaLLM):
        self.llm = llm
        self._setup_prompts()

    def _setup_prompts(self):  #currently this isn't needed as we are using the embeddings to find the answer, but it can be used later if needed
        """Setup context-aware prompts for all supported languages."""
        self.prompts = {
            'en': ChatPromptTemplate.from_template(
                """You are 'Eltorion', a helpful AI assistant for an ERP system named REST ERP.
Your task is to answer the user's general questions about the system and ERP concepts.
When answering:
- Be professional and helpful.
- Provide practical, actionable advice when possible.
- If the question needs specific data, suggest the user ask for reports or specific numbers.
- Keep responses concise but comprehensive.

User Question: {question}

Response:"""
            ),
            'ar': ChatPromptTemplate.from_template(
                """أنت 'Eltorion'، مساعد ذكاء اصطناعي متخصص في نظام 'REST ERP'.
مهمتك هي الإجابة على أسئلة المستخدم العامة حول النظام ومفاهيم تخطيط موارد المؤسسات.
عند الإجابة:
- كن محترفاً ومتعاوناً.
- قدم نصائح عملية وقابلة للتنفيذ كلما أمكن.
- إذا كان السؤال يتطلب بيانات محددة، اقترح على المستخدم أن يطلب تقارير أو أرقام معينة.
- حافظ على أن تكون إجاباتك مختصرة وشاملة.

سؤال المستخدم: {question}

الرد:"""
            ),
            'fr': ChatPromptTemplate.from_template(
                """Vous êtes 'Eltorion', un assistant IA pour le système ERP 'REST ERP'.
Votre tâche est de répondre aux questions générales de l'utilisateur sur le système et les concepts ERP.
En répondant :
- Soyez professionnel et serviable.
- Fournissez des conseils pratiques et exploitables lorsque cela est possible.
- Si la question nécessite des données spécifiques, suggérez à l'utilisateur de demander des rapports ou des chiffres précis.
- Gardez des réponses concises mais complètes.

Question de l'utilisateur : {question}

Réponse :"""
            )
        }

    def respond(self, question: str, lang_code: str = 'en') -> str:
        """
        Generate a response for a general query in the specified language. edit: find an answer by searching using the embeddings
        """
        try:
            response = ""
            
            # get the embedding for the question to search for similar chunks
            q_embedding = embedder.encode(question, show_progress_bar=False)

            #as it is only one query we need to make it in shape of (1,-1), as search method expects a 2D array
            q_embedding = q_embedding.reshape(1, -1)

            index = _load_index() #load the vectors to search for the closest matches

            # Search
            D, I = index.search(q_embedding, k=3)  #get the closest matches, to be then edited by the LLM

            #to get rid of duplicates:
            seen_indx = set()

            if I[0][0] == -1:  # no match found
                logger.warning('no relevant info in general agent search')
            else:
                for i in range(len(I[0])):
                    vec_id = I[0][i]

                    if vec_id not in seen_indx:
                        #save the index in a set to avoid duplicate values    
                        seen_indx.add(vec_id)

                        metadata = cache_service.get(f"vec_meta:{vec_id}") # Get metadata for the a match
                        if metadata:
                            response += f"\n\nText: {metadata.get('chunk', 'No relevant information found.')}"

                    # get the index of the next vector in case of splitted text
                    next_vec = vec_id + 1

                    if next_vec >= index.ntotal: #we reached the end of the vectors so, no next vec
                        continue

                    if next_vec not in seen_indx: # means that we don't have a duplicate so we can add it
                        seen_indx.add(next_vec)
                        next_metadata = cache_service.get(f"vec_meta:{next_vec}") # Get metadata of the next match in case of splitted text
                        if next_metadata:
                            response += f"\n\nText: {next_metadata.get('chunk', 'No relevant information found.')}"
            
            #in case of empty chunks
            if not response.strip():
                logger.warning('no relevant info in general agent search')
                response = 'No relevant information found.'

            return response
        
        except Exception as e:
            logger.error(f"General Agent error: {e}", exc_info=True)
            raise HTTPException(
                status_code=500,
                detail="Failed to generate response for general query"
            )