# /home/bk_anupam/code/LLM_agents/RAG_BOT/agent/agent_node.py
import json
import os
import sys
from typing import List, Optional
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import (
    HumanMessage,
    AIMessage,
    SystemMessage
)

# Add the project root to the Python path
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
sys.path.insert(0, project_root)

from RAG_BOT.config import Config
from RAG_BOT.logger import logger
from RAG_BOT.agent.state import AgentState
# Import the new prompt helper functions
from RAG_BOT.agent.prompts import get_final_answer_chat_prompt
from RAG_BOT import utils # Assuming utils has parse_json_answer


def agent_node(state: AgentState, llm: ChatGoogleGenerativeAI, llm_with_tools: ChatGoogleGenerativeAI):
    """
    Handles initial query, decides first action, and generates final response.
    Now ensures final response adheres to JSON format defined in FINAL_ANSWER_PROMPT.
    """
    logger.info("--- Executing Agent Node ---")
    messages = state['messages']
    last_message = messages[-1]
    language_code = state.get('language_code', 'en') # Default to English if not set

    # 1. Handle Initial User Query
    if isinstance(last_message, HumanMessage):
        logger.info("Handling initial user query: " + last_message.content)
        original_query_content = last_message.content
        # Decide whether to retrieve context or answer directly (usually retrieve)
        system_prompt_msg = SystemMessage(content=Config.get_system_prompt(language_code))
        # Use LLM with tools to decide if tool call is needed
        response = llm_with_tools.invoke([system_prompt_msg] + messages)
        logger.info("LLM invoked for initial decision.")
        # Update state for the first retrieval attempt
        return {
            "messages": [response],
            "original_query": original_query_content,
            "current_query": original_query_content, # Start with original query
            "retry_attempted": False,
            "evaluation_result": None, # Reset evaluation
            "context": None, # Reset context
            "language_code": language_code # Store language code
        }

    # 2. Generate Final Answer or "Cannot Find" Message (in JSON format)
    else:
        logger.info("Generating final response.")
        evaluation = state.get('evaluation_result')
        original_query = state.get('original_query')
        context = state.get('context')         

        # Check if we should generate a proper answer
        # This happens if context is sufficient OR if agent decided to answer directly
        should_generate_answer = (evaluation == 'sufficient' or(evaluation is None and not state.get('retry_attempted')))
        if should_generate_answer:
            log_context_status = "sufficient" if evaluation == 'sufficient' else "answering directly (no evaluation)"
            logger.info(f"Context {log_context_status}. Generating final answer.")
            # Use base LLM without tools for response generation
            final_answer_prompt = get_final_answer_chat_prompt(language_code) 
            logger.debug(f"Final answer prompt:  {final_answer_prompt}")
            final_answer_chain = final_answer_prompt | llm
            # Provide empty context if none was retrieved (direct answer case)
            final_answer  = final_answer_chain.invoke({      
                "system_base": Config.get_system_prompt(language_code),          
                "original_query": original_query,
                "context": context if context else "N/A" # Provide N/A if no context
            })
            if not isinstance(final_answer, AIMessage):
                final_answer = AIMessage(content=str(final_answer.content),
                                        response_metadata=getattr(final_answer, 'response_metadata', {}))
            return {"messages": [final_answer]}
        else: # Evaluation was insufficient after retry, or some other error state
            logger.info("Context insufficient after retry or error. Generating 'cannot find' message in JSON format.")
            # Format the "cannot find" message as JSON
            cannot_find_content = {
                "answer": "Relevant information cannot be found in the database to answer the question. Please try reframing your question."
            }
            # Wrap in markdown code block as per the prompt's example output
            cannot_find_json_string = f"```json\n{json.dumps(cannot_find_content, indent=2)}\n```"
            cannot_find_message = AIMessage(content=cannot_find_json_string)
            return {"messages": [cannot_find_message]}