import logging
from rag.retriever import retrieve_all
from rag.prompts import system_prompt, get_context_prompt
from services.llm import groq_client
from rag.router import classify_query

logger = logging.getLogger(__name__)

def resolve_category(question: str, doc_type: str) -> str:
    """Determine the category via the AI Router or manual override."""
    if doc_type == "Auto (AI Router)":
        classification = classify_query(question)
        logger.info(f"AI Router classified query as: '{classification}'")
        return classification
    elif doc_type == "Policy":
        return "policy"
    elif doc_type == "Product":
        return "product"
    else:
        return "none"

def get_retrieved_context(question: str, resolved_type: str) -> str:
    """Retrieve filtered context from vector store or skip it if the category is 'none'."""
    if resolved_type == "none":
        logger.info("Category is 'none'. Skipping context retrieval and RAG pipeline.")
        return ""
    
    docs = retrieve_all(question, resolved_type)
    return "\n\n".join(d.page_content for d in docs)

def build_llm_messages(question: str, context: str, history: list = None) -> list:
    """Construct prompt message structure for LLM."""
    messages = [
        {"role": "system", "content": system_prompt}
    ]

    if context:
        logger.info("Injecting retrieved document context into standard prompt.")
        context_prompt = get_context_prompt(context)
        messages.append({"role": "system", "content": context_prompt})
    else:
        logger.warning("No context found. Proceeding with system prompt only.")

    if history:
        for msg in history:
            messages.append({"role": msg["role"], "content": msg["content"]})

    messages.append({"role": "user", "content": question})
    return messages

def generate_llm_response(messages: list) -> str:
    """Interact with Groq API to retrieve completion response."""
    logger.info("Requesting chat completion from Groq LLM...")
    response = groq_client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=messages,
        max_tokens=400,
        temperature=0.6
    )
    answer = response.choices[0].message.content
    logger.info("Chat completion completed successfully.")
    return answer

def chat(question: str, doc_type: str = "Auto (AI Router)", history: list = None) -> str:
    """Coordinating function to run the full RAG chat workflow."""
    logger.info(f"Initiating chat logic for question: '{question}' (selected mode: {doc_type}, history length: {len(history) if history else 0})")
    try:
        resolved_type = resolve_category(question, doc_type)
        context = get_retrieved_context(question, resolved_type)
        messages = build_llm_messages(question, context, history)
        return generate_llm_response(messages)
    except Exception as e:
        logger.error(f"Error in modular chat workflow: {e}", exc_info=True)
        return f"An error occurred while formulating a response: {e}"