import os
from dotenv import load_dotenv
from typing import List, Dict, Any, Optional
import tempfile
import re
import json
import requests
from urllib.parse import urlparse
import pytesseract
import cmath
import pandas as pd
import uuid
import numpy as np

from PIL import Image
import pandasql as psql



### Langraph ###
from langgraph.graph import START, StateGraph, MessagesState
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.document_loaders import ArxivLoader
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_groq import ChatGroq
from langchain_huggingface import (
    ChatHuggingFace,
    HuggingFaceEndpoint,
    HuggingFaceEmbeddings,
)

from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.tools import tool
from langchain.tools.retriever import create_retriever_tool

from langchain.vectorstores import Chroma
import os
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
os.environ["GROQ_API_KEY"] = os.environ.get("GROQ_API_KEY")
os.environ["TAVILY_API_KEY"] = os.environ.get("TAVILY_API_KEY")
load_dotenv()

### =============== SERACH TOOLS =============== ###

def _format_documents_as_pseudo_xml(documents, result_tag="Document", truncate_chars=None) -> str:
    """
    Formats a list of LangChain documents into a pseudo-XML string.

    Args:
        documents (list): List of documents to format.
        result_tag (str): Tag to use for each result block.
        truncate_chars (int, optional): If set, truncates content to this many characters.

    Returns:
        str: Formatted string.
    """
    return "\n\n---\n\n".join(
        f'<{result_tag} source="{doc.metadata.get("source", "")}" page="{doc.metadata.get("page", "")}">\n'
        f'{doc.page_content[:truncate_chars] if truncate_chars else doc.page_content}\n'
        f'</{result_tag}>'
        for doc in documents
    )


@tool
def wiki_search(query: str) -> str:
    """
    Search Wikipedia and return up to 2 results formatted as pseudo-XML.

    Args:
        query (str): The search query.

    Returns:
        str: Wikipedia search results in pseudo-XML format.
    """
    query = query.strip()
    if not query:
        return "Error: Empty query provided."

    try:
        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
        return {"wiki_serach" : _format_documents_as_pseudo_xml(search_docs, result_tag="WikipediaResult")}
    except Exception as e:
        return f"Error during Wikipedia search: {str(e)}"


@tool
def web_search(query: str) -> str:
    """
    Perform a web search using Tavily and return up to 3 results formatted as pseudo-XML.

    Args:
        query (str): The search query.

    Returns:
        str: Web search results in pseudo-XML format.
    """
    query = query.strip()
    if not query:
        return "Error: Empty query provided."

    try:
        search_docs = TavilySearchResults(max_results=3).invoke(query=query)
        return {"web_search": _format_documents_as_pseudo_xml(search_docs, result_tag="WebResult")}
    except Exception as e:
        return f"Error during web search: {str(e)}"


@tool
def arxiv_search(query: str) -> str:
    """
    Search Arxiv for academic papers and return up to 3 results formatted as pseudo-XML (truncated).

    Args:
        query (str): The search query.

    Returns:
        str: Arxiv search results in pseudo-XML format (first 1000 chars of each result).
    """
    query = query.strip()
    if not query:
        return "Error: Empty query provided."

    try:
        search_docs = ArxivLoader(query=query, load_max_docs=3).load()
        return {"arxiv_search": _format_documents_as_pseudo_xml(search_docs, result_tag="ArxivResult")}
    except Exception as e:
        return f"Error during Arxiv search: {str(e)}"


### =============== MATHEMATICAL TOOLS =============== ###


@tool
def multiply(a: float, b: float) -> float:
    """
    Multiplies two numbers.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a * b


@tool
def add(a: float, b: float) -> float:
    """
    Adds two numbers.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a + b


@tool
def subtract(a: float, b: float) -> int:
    """
    Subtracts two numbers.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a - b


@tool
def divide(a: float, b: float) -> float:
    """
    Divides two numbers.
    Args:
        a (float): the first float number
        b (float): the second float number
    """
    if b == 0:
        raise ValueError("Cannot divided by zero.")
    return a / b


@tool
def modulus(a: int, b: int) -> int:
    """
    Get the modulus of two numbers.
    Args:
        a (int): the first number
        b (int): the second number
    """
    return a % b


@tool
def power(a: float, b: float) -> float:
    """
    Get the power of two numbers.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a**b


@tool
def square_root(a: float) -> float | complex:
    """
    Get the square root of a number.
    Args:
        a (float): the number to get the square root of
    """
    if a >= 0:
        return a**0.5
    return cmath.sqrt(a)


### =============== DOCUMENT PROCESSING TOOLS =============== ###


def generate_summary(df: pd.DataFrame) -> str:
    """
    Generate a text summary of a DataFrame including row/column counts and summary statistics.

    Args:
        df (pd.DataFrame): The DataFrame to summarize.

    Returns:
        str: A summary string describing the DataFrame.
    """
    summary = f"Loaded file with {len(df)} rows and {len(df.columns)} columns.\n"
    summary += f"Columns: {', '.join(df.columns)}\n\n"
    summary += "Summary statistics:\n"
    summary += str(df.describe(include='all'))
    return summary

def validate_file_path(file_path: str) -> None:
    """
    Validate whether a file path exists and is safe to use.

    Args:
        file_path (str): Path to the file.

    Raises:
        FileNotFoundError: If the file does not exist.
        ValueError: If the file is too large to process safely.
    """
    if not os.path.exists(file_path) or not os.path.isfile(file_path):
        raise FileNotFoundError(f"File does not exist: {file_path}")
    if os.path.getsize(file_path) > 100 * 1024 * 1024:  # 100MB cap
        raise ValueError("File is too large to analyze securely.")


@tool
def save_and_read_file(content: str, filename: Optional[str] = None) -> str:
    """
    Save text content to a file in the system's temporary directory.

    Args:
        content (str): The text content to save in the file.
        filename (Optional[str]): Optional name of the file. If not provided, a unique temporary filename is generated.

    Returns:
        str: File path where the content was saved, ready for further processing.
    """
    
    temp_dir = tempfile.gettempdir()
    filepath = os.path.join(temp_dir, filename or f"file_{uuid.uuid4().hex[:8]}.txt")
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(content)
    return f"File saved to {filepath}. You can read this file to process its contents."


@tool
def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
    """
    Download a file from a specified URL and save it to a temporary directory.

    Args:
        url (str): The URL pointing to the file to download.
        filename (Optional[str]): Optional filename for the saved file. If not specified, the name is inferred or generated.

    Returns:
        str: Message with the file path if download is successful, otherwise an error message.

    Raises:
        Exception: If the request fails or file cannot be written.
    """
    try:
        if not filename:
            filename = os.path.basename(urlparse(url).path) or f"downloaded_{uuid.uuid4().hex[:8]}"
        filepath = os.path.join(tempfile.gettempdir(), filename)

        response = requests.get(url, stream=True, timeout=10)
        if response.status_code != 200:
            return f"Failed to download. HTTP status: {response.status_code}"

        with open(filepath, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)

        return f"File downloaded to {filepath}. You can read this file to process its contents."
    except Exception as e:
        return f"Error downloading file: {str(e)}"


@tool
def extract_text_from_image(image_path: str) -> str:
    """
    Extract text from an image file using OCR (Optical Character Recognition).

    Args:
        image_path (str): Path to the image file.

    Returns:
        str: Extracted text from the image, or an error message if extraction fails.

    Raises:
        Exception: If the image cannot be loaded or OCR fails.
    """
    
    try:
        validate_file_path(image_path)
        image = Image.open(image_path)
        text = pytesseract.image_to_string(image)
        return f"Extracted text from image:\n\n{text}"
    except Exception as e:
        return f"Error extracting text from image: {str(e)}"


@tool
def analyze_csv_file(file_path: str, query: str = "") -> str:
    """
    Analyze a CSV file using pandas. Optionally run an SQL-like query using pandasql.

    Args:
        file_path (str): Path to the CSV file.
        query (str): Optional SQL-like query to run on the data.

    Returns:
        str: A summary of the CSV data and (if provided) query results.

    Raises:
        Exception: If file loading, parsing, or querying fails.
    """
    try:
        validate_file_path(file_path)
        df = pd.read_csv(file_path)

        result = generate_summary(df)
        if query.strip():
            query_result = psql.sqldf(query, locals())
            result += "\n\nQuery Result:\n"
            result += str(query_result.head())

        return result
    except Exception as e:
        return f"Error analyzing CSV file: {str(e)}"


@tool
def analyze_excel_file(file_path: str, query: str = "") -> str:
    """
    Analyze an Excel (.xlsx) file using pandas. Optionally run an SQL-like query using pandasql.

    Args:
        file_path (str): Path to the Excel file.
        query (str): Optional SQL-like query to run on the data.

    Returns:
        str: A summary of the Excel data and (if applicable) query results.

    Raises:
        Exception: If file loading, parsing, or querying fails.
    """
    try:
        validate_file_path(file_path)
        df = pd.read_excel(file_path)

        result = generate_summary(df)
        if query.strip():
            query_result = psql.sqldf(query, locals())
            result += "\n\nQuery Result:\n"
            result += str(query_result.head())

        return result
    except Exception as e:
        return f"Error analyzing Excel file: {str(e)}"

@tool
def youtube_transcript(video_url: str) -> str:
    """
    Extract transcript from a YouTube video.
    Args:
        video_url (str): The full URL to the YouTube video.
    Returns:
        str: The text transcript or an error.
    """
    try:
        from youtube_transcript_api import YouTubeTranscriptApi
        from urllib.parse import urlparse, parse_qs

        video_id = parse_qs(urlparse(video_url).query)['v'][0]
        transcript = YouTubeTranscriptApi.get_transcript(video_id)
        text = "\n".join([entry['text'] for entry in transcript])
        return f"<YouTubeTranscript video_id=\"{video_id}\">\n{text}\n</YouTubeTranscript>"
    except Exception as e:
        return f"Error extracting YouTube transcript: {str(e)}"


# load the system prompt from the file
with open("system_prompt.txt", "r", encoding="utf-8") as f:
    system_prompt = f.read()
print(system_prompt)

# System message
sys_msg = SystemMessage(content=system_prompt)

# build a retriever
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")

# Initialize Chroma vector store
# - persist_directory stores the collection locally, or set to None for in-memory
vector_store = Chroma(
    collection_name="documents",
    embedding_function=embeddings,
    persist_directory="./chroma_db"  # create this folder or set None for in-memory
)


create_retriever_tool = create_retriever_tool(
    retriever=vector_store.as_retriever(),
    name="Question Search",
    description="A tool to retrieve similar questions from a vector store.",
)


tools = [
    web_search,
    wiki_search,
    arxiv_search,
    multiply,
    add,
    subtract,
    divide,
    modulus,
    power,
    square_root,
    save_and_read_file,
    download_file_from_url,
    extract_text_from_image,
    analyze_csv_file,
    analyze_excel_file,
    youtube_transcript,
]


# Build graph function
def build_graph(provider: str = "groq"):
    """Build the graph"""
    # Load environment variables from .env file
    if provider == "groq":
        # Groq https://console.groq.com/docs/models
        llm = ChatGroq(model="qwen-qwq-32b", temperature=0)
    elif provider == "huggingface":
        # TODO: Add huggingface endpoint
        llm = ChatHuggingFace(
            llm=HuggingFaceEndpoint(
                repo_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
                task="text-generation",  # for chat‐style use “text-generation”
                max_new_tokens=1024,
                do_sample=False,
                repetition_penalty=1.03,
                temperature=0,
            ),
            verbose=True,
        )
    else:
        raise ValueError("Invalid provider. Choose 'groq' or 'huggingface'.")
    # Bind tools to LLM
    llm_with_tools = llm.bind_tools(tools)

    # Node
    def assistant(state: MessagesState):
        """Assistant node"""
        return {"messages": [llm_with_tools.invoke(state["messages"])]}

    def retriever(state: MessagesState):
        similar_question = vector_store.similarity_search(state["messages"][0].content)

        if similar_question:
            example_msg = HumanMessage(
                content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
            )
            return {"messages": [sys_msg] + state["messages"] + [example_msg]}
        else:
            return {"messages": [sys_msg] + state["messages"]}

    builder = StateGraph(MessagesState)
    builder.add_node("retriever", retriever)
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))
    builder.add_edge(START, "retriever")
    builder.add_edge("retriever", "assistant")
    builder.add_conditional_edges(
        "assistant",
        tools_condition,
    )
    builder.add_edge("tools", "assistant")

    # Compile graph
    return builder.compile()


# test
if __name__ == "__main__":
    question = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."
    graph = build_graph(provider="groq")
    messages = [HumanMessage(content=question)]
    messages = graph.invoke({"messages": messages})
    for m in messages["messages"]:
        m.pretty_print()