from __future__ import annotations
import asyncio
import os
import json
import logging
import time
from typing import List, Dict, Tuple, Optional, Any, Literal

from fastapi import HTTPException
from pydantic import BaseModel, Field

from s3_utils import download_chroma_folder_from_s3
import torch
import chromadb
from chromadb.api import Collection
from chromadb import PersistentClient

from modal import App, Image, Secret, fastapi_endpoint, enter, method
from dotenv import load_dotenv

load_dotenv()

logging.basicConfig(level=logging.INFO, format='{"time": "%(asctime)s", "level": "%(levelname)s", "message": "%(message)s"}')
logger = logging.getLogger(__name__)

# LLM_MODEL_GPU_ID = "meta-llama/Llama-3.1-8B-Instruct"
TINY_MODEL_ID = "meta-llama/Llama-3.2-3B-Instruct"
DEVICE = "cuda:0"
LLAMA_3_CONTEXT_WINDOW = 8192
SAFETY_BUFFER = 50

RETRIEVE_TOP_K_GPU = 8
MAX_NEW_TOKENS_GPU = 1024

CROSS_ENCODER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"

CHROMA_DIR = os.getenv("CHROMA_DIR")
CHROMA_DIR_INF = "/" + CHROMA_DIR
CHROMA_COLLECTION = os.getenv("CHROMA_COLLECTION")
CHROMA_CACHE_COLLECTION = os.getenv("CHROMA_CACHE_COLLECTION")

REQUEST_TIMEOUT_SEC = 1800

rag_image = (
    Image.from_registry("nvidia/cuda:12.1.0-base-ubuntu22.04", add_python="3.11")
    .apt_install("git")
    .pip_install_from_requirements("requirements.txt")
    .env({"HF_HOME": "/root/.cache/huggingface/hub"})
    .add_local_python_source("s3_utils", copy=True)
    .add_local_dir(
        local_path="./",
        remote_path="/usr/src/app/",
        ignore=[
            "__pycache__/", "utils/", "Dockerfile", "chroma_db_files/", "model/",
            "hg_login.py", "infer.py", "inference_chroma.py", "initial.py", "README.md",
            "requirements_heavy.txt", "requirements_light.txt", "upload_model.py", ".env",
            ".git/", "*.pyc", ".python-version", "test_*.py", "experiments/", "logs/"
        ],
        copy=True
    )
)

app = App("who-rag-llama3-gpu-api", image=rag_image)

class HistoryMessage(BaseModel):
    role: Literal['user', 'assistant']
    content: str

class QueryRequest(BaseModel):
    query: str = Field(..., description="The user's latest message.")
    history: List[HistoryMessage] = Field(default_factory=list, description="The previous turns of the conversation.")
    stream: bool = Field(False)

class RAGResponse(BaseModel):
    query: str = Field(..., description="The original user query.")
    answer: str = Field(..., description="The final answer generated by the LLM.")
    sources: List[str] = Field(..., description="Unique source URLs used for the answer.")
    context_chunks: List[str] = Field(..., description="The final context chunks (text only) sent to the LLM.")
    expanded_queries: List[str] = Field(..., description="Queries used for retrieval.")

@app.cls(
    gpu="T4",
    secrets=[
        Secret.from_name("aws-credentials"), 
        Secret.from_name("chromadb"),
        Secret.from_name("huggingface-token")
    ],
    timeout=1080,
    startup_timeout=600,
    memory=32768
)
class RagService:
    # gpu_pipeline: Any = None
    # tokenizer: Any = None
    chroma_collection: Optional[Collection] = None
    cache_collection: Optional[Collection] = None
    cross_encoder: Any = None
    embedding_model: Any = None
    intent_pipeline: Any = None
    intent_tokenizer: Any = None

    @enter()
    def setup(self):
        """Initialize all models once during container startup"""
        import torch
        from sentence_transformers import CrossEncoder
        from fastembed import TextEmbedding
        from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
        from transformers import BitsAndBytesConfig

        logger.info("Starting Modal Service setup...")

        try:
            os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
            torch.cuda.empty_cache()

            client = self._initialize_chroma_client()
            self.chroma_collection = client.get_collection(name=CHROMA_COLLECTION)
            self.cache_collection = client.get_or_create_collection(name=CHROMA_CACHE_COLLECTION)
            logger.info(f"Loaded collection: {CHROMA_COLLECTION}")

            self.embedding_model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
            _ = list(self.embedding_model.embed(["warmup"]))
            logger.info("Embedding model loaded")

            self.cross_encoder = CrossEncoder(CROSS_ENCODER_MODEL, device="cpu")
            logger.info("Cross-encoder loaded")

            logger.info(f"Loading intent model: {TINY_MODEL_ID}")
            self.intent_pipeline, self.intent_tokenizer = self._initialize_lightweight_pipeline(TINY_MODEL_ID)
            logger.info("Intent model loaded")

            torch.cuda.empty_cache()

            # self.gpu_pipeline, self.tokenizer = self._initialize_llm_pipeline(LLM_MODEL_GPU_ID)
            logger.info("Main LLM loaded")

            logger.info("All RAG components loaded successfully")

        except Exception as e:
            logger.error(f"Setup failed: {e}", exc_info=True)
            raise RuntimeError(f"Service setup failed: {e}")

    def _initialize_lightweight_pipeline(self, model_id: str):
        """Initialize lightweight pipeline for intent classification"""
        from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
        from transformers import BitsAndBytesConfig
        
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )
        
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            device_map="auto",
            trust_remote_code=True,
            quantization_config=quantization_config,
            dtype=torch.bfloat16
        )
        
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        
        if not getattr(tokenizer, "chat_template", None):
            tokenizer.chat_template = self._get_chat_template()
            
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
            
        pipe = pipeline(
            "text-generation",
            model=model,
            tokenizer=tokenizer,
            device_map="auto",
            dtype=torch.bfloat16
        )
        
        return pipe, tokenizer

    def _initialize_llm_pipeline(self, model_id: str):
        """Initialize the main LLM pipeline"""
        from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
        from transformers import BitsAndBytesConfig
        
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )
        
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            device_map="auto",
            trust_remote_code=True,
            quantization_config=quantization_config,
            dtype=torch.bfloat16
        )
        
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        
        if not getattr(tokenizer, "chat_template", None):
            tokenizer.chat_template = self._get_chat_template()
            
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
            
        pipe = pipeline(
            "text-generation",
            model=model,
            tokenizer=tokenizer,
            device_map="auto",
            dtype=torch.bfloat16
        )
        
        return pipe, tokenizer

    @staticmethod
    def _get_chat_template():
        return (
            "{% for message in messages %}"
            "{% if message['role'] == 'system' %}"
                "{{ message['content'] }} "
            "{% elif message['role'] == 'user' %}"
                "{{ '<|start_header_id|>user<|end_header_id|>\\n' + message['content'] + '<|eot_id|>' }} "
            "{% elif message['role'] == 'assistant' %}"
                "{{ '<|start_header_id|>assistant<|end_header_id|>\\n' + message['content'] + '<|eot_id|>' }} "
            "{% endif %}"
            "{% endfor %}"
            "{% if add_generation_prompt %}"
                "{{ '<|start_header_id|>assistant<|end_header_id|>\\n' }} "
            "{% endif %}"
        )

    @staticmethod
    def _initialize_chroma_client() -> chromadb.PersistentClient:
        logger.info("Starting Chroma client initialization...")
        try:
            if CHROMA_DIR is None:
                raise RuntimeError("CHROMA_DIR environment variable is not set.")
            download_chroma_folder_from_s3(CHROMA_DIR, CHROMA_DIR_INF)
            logger.info(f"Chroma data downloaded from S3 to {CHROMA_DIR_INF}.")
        except Exception as e:
            logger.error(f"Failed to download Chroma index from S3: {e}")
            raise RuntimeError("Chroma index S3 download failed.")
        
        try:
            client = PersistentClient(path=CHROMA_DIR_INF, settings=chromadb.Settings(allow_reset=False))
            logger.info("Chroma client initialized successfully.")
        except Exception as e:
            logger.error(f"Failed to load Chroma index from path: {e}")
            raise RuntimeError("Chroma index failed to load.")
        return client

    @staticmethod
    def _call_llm_pipeline(pipe: Optional[object], prompt_text: str, deterministic: bool = False, 
                          max_new_tokens: int = MAX_NEW_TOKENS_GPU, is_expansion: bool = False) -> str:
        import torch
        if pipe is None or not hasattr(pipe, "tokenizer"):
            raise HTTPException(status_code=503, detail="LLM pipeline is not available.")
            
        temp = 0.0 if deterministic else 0.1 if is_expansion else 0.6
        
        try:
            with torch.inference_mode():
                outputs = pipe(
                    prompt_text,
                    max_new_tokens=max_new_tokens,
                    temperature=(temp if temp > 0.0 else None),
                    do_sample=True if temp > 0.0 else False,
                    pad_token_id=pipe.tokenizer.eos_token_id,
                    return_full_text=False
                )
                
            if isinstance(outputs, list) and len(outputs) > 0 and isinstance(outputs[0], dict):
                text = outputs[0].get('generated_text', "")
            elif isinstance(outputs, dict):
                text = outputs.get('generated_text', "")
            else:
                text = str(outputs)
                
            text = text.strip()
            for token in ['<|eot_id|>', '<|end_of_text|>']:
                if token in text:
                    text = text.split(token)[0].strip()
            return text
            
        except Exception as e:
            logger.error(f"Error calling LLM pipeline: {e}", exc_info=True)
            raise HTTPException(status_code=500, detail=f"LLM generation failed: {str(e)}")

    def _build_prompt(self, user_query: str, context: List[Dict], summary: str) -> List[Dict]:
        context_text = "\n---\n".join([f"Source: {c.get('url', 'N/A')}\nChunk: {c['text']}" for c in context]) if context else "No relevant context found."
        
        system_prompt = (
            "You are a helpful and harmless medical assistant, specialized in answering health-related questions "
            "based ONLY on the provided retrieved context. Follow these strict rules:\n"
            "1. **DO NOT** use any external knowledge. If the answer is not in the context, state that you cannot find "
            "the information in the knowledge base.\n"
            "2. Cite your sources using the URL/Source ID provided in the context (e.g., [Source: URL]). Do not generate fake URLs.\n"
            "3. If the user's query is purely conversational, greet them or respond appropriately without referencing the context.\n"
        )
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "system", "content": f"PREVIOUS CONVERSATION SUMMARY: {summary}" if summary else "PREVIOUS CONVERSATION SUMMARY: None"},
            {"role": "system", "content": f"RETRIEVED CONTEXT:\n{context_text}"},
            {"role": "user", "content": user_query}
        ]
        return messages

    def _get_token_count(self, msg_list: List[Dict]) -> int:
        if not self.intent_tokenizer:
            return 0
        prompt_text = self.intent_tokenizer.apply_chat_template(msg_list, tokenize=False, add_generation_prompt=True)
        return len(self.intent_tokenizer.encode(prompt_text, add_special_tokens=False))

    @method()
    async def classify_intent(self, query: str) -> str:
        """Classify query intent using the pre-loaded intent pipeline"""
        if not self.intent_pipeline or not self.intent_tokenizer:
            raise HTTPException(status_code=503, detail="Intent classification model not available")
            
        system_prompt = """You are a query classification robot. You MUST respond with ONLY ONE JSON object:
        {"intent": "MEDICAL"}
        {"intent": "GREET"}
        {"intent": "OFF_TOPIC"}
        {"intent": "HARMFUL"}
        """
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": "Query: What are the symptoms of COVID-19?"},
            {"role": "assistant", "content": '{"intent": "MEDICAL"}'},
            {"role": "user", "content": f"Query: {query}"}
        ]
        
        prompt_text = self.intent_tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        
        try:
            llm_output = await self._run_with_timeout(
                asyncio.to_thread(
                    self._call_llm_pipeline, 
                    self.intent_pipeline, 
                    prompt_text, 
                    True, 25, False
                ),
                # timeout_seconds=30,
                timeout_message="Intent classification timed out"
            )
            
            clean_output = llm_output.strip().replace("```json", "").replace("```", "")
            start_idx = clean_output.find('{')
            end_idx = clean_output.rfind('}')
            if start_idx != -1 and end_idx != -1:
                json_str = clean_output[start_idx: end_idx + 1]
                data = json.loads(json_str)
                return data.get("intent", "UNKNOWN")
                
        except Exception as e:
            logger.error(f"Failed to parse JSON classifier output: {e}. Raw: {llm_output}")
            
        return self._rule_based_intent_classification(query)

    def _rule_based_intent_classification(self, query: str) -> str:
        """Fallback rule-based intent classification"""
        query_lower = query.lower().strip()
        
        greeting_words = ['hello', 'hi', 'hey', 'greetings', 'good morning', 'good afternoon', 'how are you']
        harmful_keywords = ['harm', 'hurt', 'kill', 'danger', 'illegal', 'prescription without', 'suicide']
        medical_keywords = ['covid', 'fever', 'pain', 'symptom', 'treatment', 'medicine', 'doctor', 'health', 'disease', 'virus']
        
        if any(word in query_lower for word in greeting_words) or len(query_lower.split()) <= 2:
            return 'GREET'
        elif any(word in query_lower for word in harmful_keywords):
            return 'HARMFUL'
        elif not any(word in query_lower for word in medical_keywords) and len(query_lower.split()) > 3:
            return 'OFF_TOPIC'
        else:
            return 'MEDICAL'

    @method()
    async def Greet(self, query: str) -> RAGResponse:
        """Handle greeting queries"""
        messages = [
            {"role": "system", "content": "You are a greeter. Respond politely to the user greeting in a single line."},
            {"role": "user", "content": query}
        ]
        
        prompt_text = self.intent_tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        
        answer = await self._run_with_timeout(
            asyncio.to_thread(self._call_llm_pipeline, self.intent_pipeline, prompt_text, True, 50, True),
            # timeout_seconds=30,
            timeout_message="Greeting response timed out"
        )
        
        return RAGResponse(
            query=query, 
            answer=answer, 
            sources=[], 
            context_chunks=[], 
            expanded_queries=[]
        )

    @method()
    async def HarmOff(self, query: str) -> RAGResponse:
        """Handle harmful/off-topic queries"""
        messages = [
            {"role": "system", "content": "You are an intelligent assistant. Inform the user that you cannot answer harmful/off-topic questions. Keep it short and brief, in one sentence."},
            {"role": "user", "content": query}
        ]
        
        prompt_text = self.intent_tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        
        answer = await self._run_with_timeout(
            asyncio.to_thread(self._call_llm_pipeline, self.intent_pipeline, prompt_text, True, 50, True),
            # timeout_seconds=30,
            timeout_message="Safety response timed out"
        )
        
        return RAGResponse(
            query=query, 
            answer=answer, 
            sources=[], 
            context_chunks=[], 
            expanded_queries=[]
        )

    @method()
    async def summarize_history(self, history: List[HistoryMessage]) -> str:
        """Summarize conversation history"""
        if not history:
            return ''
            
        history_text = "\n".join([f"{h.role}: {h.content}" for h in history[-8:]])
        summarizer_prompt = f"You are an intelligent agent who summarizes conversations. Your summary should be concise, coherent, and focus on the main topic and specific entities discussed.\nCONVERSATION HISTORY:\n{history_text}\nCONCISE SUMMARY:\n"
        
        messages = [{"role": "system", "content": summarizer_prompt}]
        prompt_text = self.intent_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        
        summary = await self._run_with_timeout(
            asyncio.to_thread(self._call_llm_pipeline, self.intent_pipeline, prompt_text, True, 150, False),
            # timeout_seconds=60,
            timeout_message="Summarization timed out"
        )
        return summary

    @method()
    async def expand_query_with_llm(self, user_query: str, summary: str, history: List[HistoryMessage]) -> List[str]:
        """Expand query for better retrieval"""
        if not history or len(history) == 0:
            expansion_prompt = f"You are a specialized query expansion engine. Generate 3 alternative, highly effective search queries to find documents relevant to the User Query. Only output the queries, one per line. Do not include the original query or any explanations.\nUser Query: {user_query}\nExpanded Queries:\n"
        else:
            history_text = "\n".join([f"{h.role}: {h.content}" for h in history])
            expansion_prompt = f"Given the conversation summary and history below, rewrite the user's latest query into a standalone, complete, and specific search query that incorporates the context of the conversation. Output only the single rewritten query.\nConversation Summary: {summary}\nConversation History:\n{history_text}\nUser's Latest Query: {user_query}\nRewritten Search Query:\n"
        
        messages = [{"role": "system", "content": expansion_prompt}]
        prompt_text = self.intent_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        
        llm_output = await self._run_with_timeout(
            asyncio.to_thread(self._call_llm_pipeline, self.intent_pipeline, prompt_text, True, 150, True),
            # timeout_seconds=60,
            timeout_message="Query expansion timed out"
        )
        
        if not history or len(history) == 0:
            expanded_queries = [q.strip() for q in llm_output.split('\n') if q.strip()]
        else:
            expanded_queries = [llm_output.strip()]
        
        expanded_queries.append(user_query)
        seen = set()
        deduped = []
        for q in expanded_queries:
            if q not in seen:
                seen.add(q)
                deduped.append(q)
        return deduped

    def retrieve_context(self, queries: List[str]) -> Tuple[List[Dict], List[str]]:
        """Retrieve context from ChromaDB"""
        if self.embedding_model is None:
            raise HTTPException(status_code=503, detail="Embedding model not loaded.")
            
        embeddings_list = [[float(x) for x in emb] for emb in self.embedding_model.embed(queries, batch_size=8)]
        results = self.chroma_collection.query(
            query_embeddings=embeddings_list,
            n_results=max(10, RETRIEVE_TOP_K_GPU * len(queries)),
            include=['documents', 'metadatas']
        )
        
        context_data = []
        source_urls = set()
        if results.get("documents") and results.get("metadatas"):
            for docs_list, metadatas_list in zip(results["documents"], results["metadatas"]):
                for doc, metadata in zip(docs_list, metadatas_list):
                    if doc and metadata:
                        context_data.append({'text': doc, 'url': metadata.get('source')})
                        if metadata.get("source"):
                            source_urls.add(metadata.get('source'))
        return context_data, list(source_urls)

    def rerank_documents(self, query: str, context: List[Dict], top_k: int) -> List[Dict]:
        """Rerank documents using cross-encoder"""
        if not context or self.cross_encoder is None:
            return context[:top_k]
            
        pairs = [(query, doc['text']) for doc in context]
        scores = self.cross_encoder.predict(pairs)
        for doc, score in zip(context, scores):
            doc['score'] = float(score)
        ranked_docs = sorted(context, key=lambda x: x['score'], reverse=True)
        return ranked_docs[:top_k]

    @method()
    async def prune_messages_to_fit_context(self, messages: List[Dict], final_context: List[Dict], summary: str, max_input_tokens: int) -> Tuple[List[Dict], List[Dict], int]:
        """Prune messages to fit within token limit"""
        if not self.intent_tokenizer:
            raise ValueError("Tokenizer not initialized for pruning.")
            
        current_context = final_context[:]
        current_summary = summary
        base_user_query = messages[-1]["content"]
        
        current_messages = self._build_prompt(base_user_query, current_context, current_summary)
        token_count = self._get_token_count(current_messages)
        
        if token_count <= max_input_tokens:
            tok_length = max_input_tokens - token_count
            return current_messages, current_context, tok_length
            
        logger.warning(f"Initial token count ({token_count}) exceeds max input ({max_input_tokens}). Starting pruning.")
        
        while token_count > max_input_tokens and current_context:
            current_context.pop()
            current_messages = self._build_prompt(base_user_query, current_context, current_summary)
            token_count = self._get_token_count(current_messages)
            
        if token_count <= max_input_tokens:
            tok_length = max_input_tokens - token_count
            return current_messages, current_context, tok_length
            
        if current_summary:
            logger.warning("Clearing conversation summary as last-ditch effort.")
            current_summary = ""
            current_messages = self._build_prompt(base_user_query, current_context, current_summary)
            token_count = self._get_token_count(current_messages)
            
        if token_count <= max_input_tokens:
            tok_length = max_input_tokens - token_count
            return current_messages, current_context, tok_length
            
        logger.error(f"Pruning failed. Even minimal prompt exceeds token limit: {token_count}. Returning empty context.")
        current_context = []
        current_messages = self._build_prompt(base_user_query, current_context, "")
        token_count = self._get_token_count(current_messages)
        tok_length = max_input_tokens - token_count if token_count < max_input_tokens else 0
        
        return current_messages, current_context, tok_length

    @fastapi_endpoint(method="POST")
    async def rag_endpoint(self, request_data: Dict[str, Any]):
        """Main RAG endpoint"""
        try:
            request = QueryRequest(**request_data)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid request format: {str(e)}")

        start = time.time()
        
        try:
            logger.info(f'Processing query: {request.query[:100]}...')
            
            intent = await self.classify_intent.remote.aio(request.query)
            logger.info(f"Intent classified as: {intent}")

            if intent == 'GREET':
                response = await self.Greet.remote.aio(request.query)
            elif intent in ["HARMFUL", "OFF_TOPIC"]:
                response = await self.HarmOff.remote.aio(request.query)
            else:
                logger.info("Starting full RAG pipeline for medical query")
                
                summary = await self.summarize_history.remote.aio(request.history)
                logger.info("History summarized")
                
                expanded_queries = await self.expand_query_with_llm.remote.aio(request.query, summary, request.history)
                logger.info(f"Expanded queries: {expanded_queries}")
                
                context_data, _ = await self._run_with_timeout(
                    asyncio.to_thread(self.retrieve_context, expanded_queries),
                    timeout_message="Document retrieval timed out"
                )
                logger.info(f"Retrieved {len(context_data)} context chunks")
                
                final_context = await self._run_with_timeout(
                    asyncio.to_thread(self.rerank_documents, request.query, context_data, RETRIEVE_TOP_K_GPU),
                    # timeout_seconds=60,
                    timeout_message="Document reranking timed out"
                )
                logger.info(f"Reranked to {len(final_context)} chunks")
                
                final_sources = list({c.get('url') for c in final_context if c.get('url')})
                
                if not final_context:
                    final_answer = "I could not find relevant documents in the knowledge base to answer your question. I can help you if you have another question."
                    context_chunks_text = []
                else:
                    initial_messages = self._build_prompt(request.query, final_context, summary)
                    max_input_tokens = LLAMA_3_CONTEXT_WINDOW - MAX_NEW_TOKENS_GPU - SAFETY_BUFFER
                    
                    final_messages, final_context_pruned, tok_length = await self.prune_messages_to_fit_context.remote.aio(
                        initial_messages, final_context, summary, max_input_tokens
                    )
                    
                    context_chunks_text = [c['text'] for c in final_context_pruned]
                    prompt_text = self.intent_tokenizer.apply_chat_template(final_messages, tokenize=False, add_generation_prompt=True)
                    
                    max_new = max(MAX_NEW_TOKENS_GPU, tok_length if isinstance(tok_length, int) and tok_length > 0 else MAX_NEW_TOKENS_GPU)
                    
                    final_answer = await self._run_with_timeout(
                        asyncio.to_thread(self._call_llm_pipeline, self.intent_pipeline, prompt_text, False, max_new, False),
                        # timeout_seconds=120,
                        timeout_message="Answer generation timed out"
                    )
                    logger.info("Generated final answer")

                response = RAGResponse(
                    query=request.query,
                    answer=final_answer,
                    sources=final_sources,
                    context_chunks=context_chunks_text,
                    expanded_queries=expanded_queries
                )

            end_time = time.time()
            logger.info(f"Total Latency: {round(end_time - start, 2)}s")
            return response.model_dump()
            
        except HTTPException:
            raise
        except Exception as e:
            logger.error(f"Unhandled exception in RAG handler: {e}", exc_info=True)
            raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")

    async def _run_with_timeout(self, awaitable: Any, timeout_seconds: int = 300, timeout_message: str = "Request timed out") -> Any:
        try:
            return await asyncio.wait_for(awaitable, timeout=timeout_seconds)
        except asyncio.TimeoutError:
            logger.warning(f"Operation timed out after {timeout_seconds}s: {timeout_message}")
            raise HTTPException(status_code=504, detail=timeout_message)
        except HTTPException:
            raise
        except Exception as e:
            logger.error(f"Unexpected error in _run_with_timeout: {e}")
            raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")