import logging
import os
import base64
import pandas as pd
import io
import contextlib

from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_community.document_loaders import AssemblyAIAudioTranscriptLoader
from langchain_community.tools import WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper
from langchain_core.tools import tool
from langchain_google_genai import ChatGoogleGenerativeAI

logger = logging.getLogger("eval_logger")


try:
    tools_llm = ChatOpenAI(model="gpt-4o", temperature=0)
except Exception as e:
    logger.error(f"Failed to initialize tools_llm (OpenAI gpt-4o) in tools.py: {e}. Ensure OPENAI_API_KEY is set.", exc_info=True)
    tools_llm = None


GEMINI_SHARED_MODEL_NAME = "gemini-2.5-pro-preview-05-06"
try:
    gemini_llm = ChatGoogleGenerativeAI(
        model=GEMINI_SHARED_MODEL_NAME, 
        temperature=0,
        timeout=360 # 6-minute timeout
    )
    logger.info(f"Successfully initialized shared Gemini model: {GEMINI_SHARED_MODEL_NAME} with a 360s timeout.")
except Exception as e:
    logger.error(f"Failed to initialize shared_gemini_llm in tools.py (model: {GEMINI_SHARED_MODEL_NAME}): {e}. Ensure GOOGLE_API_KEY is set and valid, and the model name is correct/available.", exc_info=True)
    gemini_llm = None


try:
    tools_llm = ChatOpenAI(model="gpt-4o", temperature=0)
except Exception as e:
    logger.error(f"Failed to initialize tools_llm in tools.py (ChatOpenAI with gpt-4o): {e}. Ensure OPENAI_API_KEY is set and .env is loaded.", exc_info=True)
    tools_llm = None # Set to None so the app can still load, but tool will fail

@tool
def analyse_image(img_path: str, question: str) -> str:
    """
    Analyses a **locally stored** image file to answer a specific question using a multimodal model.
    IMPORTANT: This tool expects a local file path for 'img_path' and cannot process web URLs directly.
    Args:
        img_path: Local path to the image file (e.g., /path/to/your/image.png).
        question: The question the user is trying to answer by analysing this image.
    Returns:
        A string containing the relevant information extracted from the image to answer the question, 
        or an error message if analysis fails.
    """
    if not tools_llm:
        return "Error: Vision LLM (gpt-4o) not initialized in tools.py. Cannot analyse image."
    if not os.path.exists(img_path):
        # This check is more critical now that we emphasize local paths.
        return f"Error: Image file not found at local path: {img_path}. This tool requires a local file path."
    
    logger.info(f"Attempting to analyse image: {img_path} for question: '{question}'")
    try:
        with open(img_path, "rb") as image_file:
            image_bytes = image_file.read()
        image_base64 = base64.b64encode(image_bytes).decode("utf-8")

        image_type = os.path.splitext(img_path)[1].lower()
        if image_type == '.jpg':
            image_type = '.jpeg'
        if image_type not in ['.png', '.jpeg', '.gif', '.webp']:
             return f"Error: Unsupported image type '{image_type}' for gpt-4o vision. Supported: PNG, JPEG, GIF, WEBP."

        prompt_text = f"Analyse this image to answer the following question: '{question}'. Focus on extracting only the information directly relevant to this question. Return only the extracted information, with no additional explanations or commentary."
        message = HumanMessage(
            content=[
                {"type": "text", "text": prompt_text},
                {"type": "image_url", "image_url": {"url": f"data:image/{image_type[1:]};base64,{image_base64}"}},
            ]
        )
        response = tools_llm.invoke([message])
        extracted_text = response.content
        logger.info(f"Successfully analysed {img_path} for question '{question}'. Response length: {len(extracted_text)}")
        return extracted_text.strip()
    except Exception as e:
        logger.error(f"Error analysing image {img_path} for question '{question}': {e}", exc_info=True)
        return f"Error during image analysis for question '{question}': {str(e)}"

@tool
def analyse_audio(audio_path: str, question: str) -> str:
    """
    Transcribes a **locally stored** audio file using AssemblyAI and then analyses the transcript 
    with a multimodal model (gpt-4o) to answer a specific question.
    IMPORTANT: This tool expects a local file path for 'audio_path' (e.g., /path/to/your/audio.mp3) 
    and **cannot process web URLs (like YouTube links) directly.**
    Args:
        audio_path: Local path to the audio file (e.g., /path/to/your/audio.mp3).
        question: The question the user is trying to answer by analysing this audio.
    Returns:
        A string containing the relevant information extracted from the audio to answer the question, 
        or an error message if analysis fails.
    """
    logger.info(f"Attempting to analyse audio from local path: {audio_path} for question: '{question}'")
    if not tools_llm: # vision_llm (gpt-4o) is used for the Q&A part
        return "Error: LLM (gpt-4o) for Q&A not initialized in tools.py. Cannot analyse audio transcript."
    if not audio_path: 
        return "Error: Audio file path not provided."
    if not os.path.exists(audio_path):
        return f"Error: Audio file not found at local path: {audio_path}. This tool requires a local file path."

    try:
        logger.info(f"Loading/transcribing audio from local file: {audio_path} using AssemblyAI.")
        loader = AssemblyAIAudioTranscriptLoader(file_path=audio_path) # AssemblyAI loader primarily works with local paths for reliability.
        docs = loader.load()
        
        if not docs or not docs[0].page_content:
            logger.error(f"AssemblyAI transcription failed or returned empty for {audio_path}.")
            return f"Error: Transcription failed or returned empty content for {audio_path}."
            
        transcript = docs[0].page_content
        logger.info(f"Successfully transcribed audio from {audio_path}. Transcript length: {len(transcript)}")
        
        qa_prompt_text = (
            f"The following is a transcript of an audio file: \n\nTranscript:\n{transcript}\n\n---\n\n"
            f"Based SOLELY on the information in the transcript above, answer the following question: '{question}'. "
            f"Provide only the direct answer as extracted or inferred from the transcript, with no additional commentary."
        )
        
        message = HumanMessage(content=qa_prompt_text)
        response = tools_llm.invoke([message])
        answer = response.content
        logger.info(f"Successfully analysed transcript from {audio_path} for question '{question}'. Answer length: {len(answer)}")
        return answer.strip()

    except Exception as e:
        logger.error(f"Error analysing audio {audio_path} for question '{question}': {e}", exc_info=True)
        if "api key" in str(e).lower() or "authenticate" in str(e).lower():
            return f"Error during audio analysis: AssemblyAI authentication failed. Please check your ASSEMBLYAI_API_KEY. Original error: {str(e)}"
        return f"Error during audio analysis for question '{question}': {str(e)}"

@tool
def execute_python_code_from_file(file_path: str, question: str) -> str:
    """
    Reads the content of a **locally stored** Python file and uses a powerful LLM (gpt-4o) 
    to answer a specific question about the Python code (e.g., its output, functionality, or errors).
    IMPORTANT: This tool expects a local file path for 'file_path' and cannot process web URLs directly.
    It does NOT actually execute the code, but rather analyses it textually.
    Args:
        file_path: Local path to the Python file (e.g., /path/to/your/script.py).
        question: The question the user is trying to answer about this Python code.
    Returns:
        A string containing the LLM's analysis or answer about the Python code, or an error message.
    """
    logger.info(f"Attempting to analyse Python file: {file_path} for question: '{question}'")
    if not tools_llm: # vision_llm (gpt-4o) is used for the analysis
        return "Error: LLM (gpt-4o) for code analysis not initialized in tools.py."
    if not file_path:
        return "Error: Python file path not provided."
    if not os.path.exists(file_path):
        return f"Error: Python file not found at local path: {file_path}. This tool requires a local file path."
    if not file_path.lower().endswith('.py'):
        return f"Error: File at {file_path} is not a Python (.py) file."

    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            python_code_content = f.read()
        
        logger.info(f"Successfully read Python file {file_path}. Content length: {len(python_code_content)}")

        analysis_prompt_text = (
            f"The following is the content of a Python file: \n\nPython Code:\n```python\n{python_code_content}\n```\n\n---\n\n"
            f"Based SOLELY on the Python code provided above, answer the following question: '{question}'. "
            f"If the question asks for the output, predict the output. If it asks about functionality, describe it. "
            f"Provide only the direct answer or analysis, with no additional commentary or explanations unless the question asks for it."
        )
        
        message = HumanMessage(content=analysis_prompt_text)
        response = tools_llm.invoke([message]) # Using vision_llm (gpt-4o) for this analysis
        answer = response.content
        logger.info(f"Successfully analysed Python code from {file_path} for question '{question}'. Answer length: {len(answer)}")
        return answer.strip()

    except Exception as e:
        logger.error(f"Error analysing Python file {file_path} for question '{question}': {e}", exc_info=True)
        return f"Error during Python file analysis for question '{question}': {str(e)}"

@tool
def execute_pandas_script_for_excel(excel_file_path: str, python_code: str) -> str:
    """
    Executes a given Python script (which should use pandas) to perform analysis on an Excel file.
    The script MUST load the Excel file using the provided 'excel_file_path' variable.
    The script MUST print its final answer to standard output. The print output will be returned as the result.
    This tool is for calculations, data manipulation, and specific lookups within the Excel file.

    Args:
        excel_file_path: The path to the Excel file that the script will process.
        python_code: A string containing the Python script to execute. 
                     Example: 
                     '''
                     import pandas as pd
                     df = pd.read_excel(excel_file_path, sheet_name=0)
                     # Perform analysis ...
                     final_answer = df["SomeColumn"].sum() # Example operation
                     print(final_answer)
                     '''
    Returns:
        The standard output from the executed script (which should be the answer), or an error message if execution fails.
    """
    logger.info(f"Attempting to execute pandas script for Excel file: {excel_file_path}")
    logger.debug(f"Python code to execute:\n{python_code}")

    if not os.path.exists(excel_file_path):
        return f"Error: Excel file not found at {excel_file_path}"

    # Prepare the local namespace for exec, including pandas and the file path
    local_namespace = {
        "pd": pd, 
        "excel_file_path": excel_file_path,
        "__builtins__": __builtins__ # Ensure basic builtins are available
    }
    
    # Capture stdout
    stdout_capture = io.StringIO()
    try:
        with contextlib.redirect_stdout(stdout_capture):
            exec(python_code, {"__builtins__": __builtins__}, local_namespace) # Provide pandas in globals, path in locals
        output = stdout_capture.getvalue().strip()
        logger.info(f"Successfully executed pandas script. Output: '{output}'")
        if not output: # If the script printed nothing, it might indicate an issue or missing print().
            return "Script executed successfully but produced no output. Ensure the script prints the final answer."
        return output
    except Exception as e:
        logger.error(f"Error executing pandas script: {e}", exc_info=True)
        # Provide a more detailed error message back to the LLM
        import traceback
        tb_str = traceback.format_exc()
        return f"Error during script execution: {str(e)}\nTraceback:\n{tb_str}"

@tool
def analyse_youtube(youtube_url: str, question: str) -> str:
    """
    Analyzes a YouTube video to answer a specific question.
    This tool is intended for questions that require understanding the visual content of the video.
    It sends the YouTube URL directly to the shared Gemini model.

    Args:
        youtube_url: The full URL of the YouTube video (e.g., https://www.youtube.com/watch?v=...). 
        question: The question to answer based on the video's content.
    Returns:
        A string containing the answer from the shared Gemini model, or an error message if analysis fails.
    """
    logger.info(f"Attempting to analyse YouTube video: {youtube_url} with shared Gemini model ({GEMINI_SHARED_MODEL_NAME}) for question: '{question}'")
    if not gemini_llm:
        return f"Error: Shared Gemini LLM ({GEMINI_SHARED_MODEL_NAME}) not initialized in tools.py. Cannot analyse YouTube video."

    try:
        
        prompt = f"Video URL: {youtube_url}\n\nQuestion: {question}\n\nBased on the video at the URL, please provide the answer."
        message = HumanMessage(content=prompt)

        response = gemini_llm.invoke([message])
        answer = response.content
        logger.info(f"Successfully analysed YouTube video {youtube_url} with shared Gemini. Answer: {answer[:200]}...")
        return answer.strip()
    except Exception as e:
        logger.error(f"Error analysing YouTube video {youtube_url} with shared Gemini ({GEMINI_SHARED_MODEL_NAME}): {e}", exc_info=True)
        return f"Error during YouTube video analysis with shared Gemini: {str(e)}"

@tool
def deep_analysis_with_gemini(question: str) -> str:
    """
    Performs a deep analysis of a complex question using a powerful shared Gemini model.
    Use this tool for questions that are multifaceted, require deep reasoning, 
    or for historical queries where standard search tools might be insufficient after initial attempts.
    This tool directly passes the question to a shared Gemini model for a comprehensive answer.

    Args:
        question: The complex question to be analyzed.
    Returns:
        A string containing the detailed answer from the shared Gemini model, or an error message.
    """
    logger.info(f"Attempting deep analysis with shared Gemini model ({GEMINI_SHARED_MODEL_NAME}) for question: '{question}'")
    if not gemini_llm:
        return f"Error: Shared Gemini LLM ({GEMINI_SHARED_MODEL_NAME}) not initialized in tools.py."

    try:
        message = HumanMessage(content=question)
        response = gemini_llm.invoke([message])
        answer = response.content
        logger.info(f"Successfully performed deep analysis with shared Gemini. Answer length: {len(answer)}")
        return answer.strip()
    except Exception as e:
        logger.error(f"Error during deep analysis with shared Gemini ({GEMINI_SHARED_MODEL_NAME}): {e}", exc_info=True)
        return f"Error during deep analysis with shared Gemini: {str(e)}"

# Initialize other tools
search_tool = DuckDuckGoSearchRun()

wikipedia_tool = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())

TOOLS = [
    analyse_image,
    search_tool,
    analyse_audio,
    execute_python_code_from_file,
    wikipedia_tool,
    execute_pandas_script_for_excel,
    analyse_youtube,
    deep_analysis_with_gemini,
]

logger.info(f"Tools initialized in tools.py: {[tool.name if hasattr(tool, 'name') else tool.__name__ for tool in TOOLS]}")