# First we initialize the model we want to use.
from langchain_anthropic import ChatAnthropic
from langchain_openai import ChatOpenAI
from litellm import image_generation

from botrun_flow_lang.langgraph_agents.agents.util.img_util import analyze_imgs
from botrun_flow_lang.langgraph_agents.agents.util.local_files import (
    upload_and_get_tmp_public_url,
)
from botrun_flow_lang.langgraph_agents.agents.util.pdf_analyzer import analyze_pdf
from botrun_flow_lang.langgraph_agents.agents.util.perplexity_search import (
    respond_with_perplexity_search,
)
from botrun_flow_lang.models.nodes.utils import scrape_single_url
from botrun_flow_lang.models.nodes.vertex_ai_search_node import VertexAISearch
from datetime import datetime
from botrun_flow_lang.langgraph_agents.agents.search_agent_graph import format_dates
from langgraph.checkpoint.memory import MemorySaver
import pytz
import asyncio
import os
from botrun_flow_lang.langgraph_agents.agents.util.plotly_util import (
    generate_plotly_files,
)
from botrun_flow_lang.langgraph_agents.agents.util.mermaid_util import (
    generate_mermaid_files,
)


# model = ChatOpenAI(model="gpt-4o", temperature=0)
model = ChatAnthropic(model="claude-3-5-sonnet-latest", temperature=0)


# For this tutorial we will use custom tool that returns pre-defined values for weather in two cities (NYC & SF)

from typing import Literal

from langchain_core.tools import tool


@tool
def get_weather(city: Literal["nyc", "sf"]):
    """Use this to get weather information."""
    if city == "nyc":
        return "It might be cloudy in nyc"
    elif city == "sf":
        return "It's always sunny in sf"
    else:
        raise AssertionError("Unknown city")


@tool
def multiply(a: int, b: int) -> int:
    """Multiply a and b.

    Args:
        a: first int
        b: second int
    """
    return a * b


# This will be a tool
@tool
def add(a: int, b: int) -> int:
    """Adds a and b.

    Args:
        a: first int
        b: second int
    """
    return a + b


@tool
def divide(a: int, b: int) -> float:
    """Divide a and b.

    Args:
        a: first int
        b: second int
    """
    return a / b


@tool
def search(keywords: str):
    """
    Use this to search the web.

    Args:
        keywords: the keywords to search for, use space to separate multiple keywords, e.g. "台灣 政府 福利"
    """
    try:
        vertex_ai_search = VertexAISearch()
        search_results = vertex_ai_search.vertex_search(
            project_id="scoop-386004",
            location="global",
            data_store_id="tw-gov-welfare_1730944342934",
            search_query=keywords,
        )
        return search_results
    except Exception as e:
        return f"Error: {e}"


@tool
def scrape(url: str):
    """
    Use this to scrape the web.
    as it provides better results for video content.

    Args:
        url: the url to scrape
    """
    try:
        return asyncio.run(scrape_single_url(url))
    except Exception as e:
        return f"Error: {e}"


@tool
def current_time():
    """
    Use this to get the current time in local timezone.
    """
    try:
        local_tz = pytz.timezone("Asia/Taipei")
        local_time = datetime.now(local_tz)
        return local_time.strftime("%Y-%m-%d %H:%M:%S %Z")
    except Exception as e:
        return f"Error: {e}"


@tool
def days_between(start_date: str, end_date: str):
    """
    Use this to get the days between two dates.

    Args:
        start_date: the start date, format: YYYY-MM-DD
        end_date: the end date, format: YYYY-MM-DD
    """
    return (
        datetime.strptime(end_date, "%Y-%m-%d")
        - datetime.strptime(start_date, "%Y-%m-%d")
    ).days


@tool
def chat_with_pdf(pdf_url: str, user_input: str):
    """
    Use this to chat with a PDF file.
    User can ask about any text, pictures, charts, and tables in PDFs that is provided. Some sample use cases:
    - Analyzing financial reports and understanding charts/tables
    - Extracting key information from legal documents
    - Translation assistance for documents
    - Converting document information into structured formats

    Data Visualization Integration:
    When the user's input indicates a need for comparison or data visualization (e.g., "compare the quarterly profits",
    "show the trend of sales"), this function can return data in a format suitable for Plotly visualization.
    The returned data will be a dictionary with a special key "__plotly_data__" containing:
    {
        "__plotly_data__": {
            "data": [...],  # Plotly data array
            "layout": {...}  # Plotly layout object
        }
    }
    You can then pass this data to create_plotly_chart to generate an interactive chart.

    If you have a local PDF file, you can use generate_tmp_public_url tool to get a public URL first:
    1. Call generate_tmp_public_url with your local PDF file path
    2. Use the returned URL as the pdf_url parameter for this function

    Args:
        pdf_url: the URL to the PDF file (can be generated using generate_tmp_public_url for local files)
        user_input: the user's input

    Returns:
        str: Analysis result or Plotly-compatible data structure if visualization is needed
    """
    print("chat_with_pdf============>", pdf_url, user_input)
    return analyze_pdf(pdf_url, user_input)


@tool
def chat_with_imgs(img_urls: list[str], user_input: str):
    """
    Use this to analyze and understand multiple images using Claude's vision capabilities.

    If you have local image files, you can use generate_tmp_public_url tool to get public URLs first:
    1. Call generate_tmp_public_url for each local image file
    2. Collect all returned URLs into a list
    3. Use the list of URLs as the img_urls parameter for this function

    Data Visualization Integration:
    When the user's input indicates a need for comparison or data visualization (e.g., "compare the values in these charts",
    "extract and plot the data from these images"), this function can return data in a format suitable for Plotly visualization.
    The returned data will be a dictionary with a special key "__plotly_data__" containing:
    {
        "__plotly_data__": {
            "data": [...],  # Plotly data array
            "layout": {...}  # Plotly layout object
        }
    }
    You can then pass this data to create_plotly_chart to generate an interactive chart.

    Supported image formats:
    - JPEG, PNG, GIF, WebP
    - Maximum file size: 5MB per image
    - Recommended size: No more than 1568 pixels in either dimension
    - Very small images (under 200 pixels) may degrade performance
    - Can analyze up to 20 images per request

    Capabilities:
    - Analyzing charts, graphs, and diagrams
    - Reading and understanding text in images
    - Describing visual content and scenes
    - Comparing multiple images in a single request
    - Answering questions about image details
    - Identifying relationships between images
    - Extracting data from charts for visualization

    Limitations:
    - Cannot identify or name specific people
    - May have reduced accuracy with low-quality or very small images
    - Limited spatial reasoning abilities
    - Cannot verify if images are AI-generated
    - Not designed for medical diagnosis or healthcare applications

    Args:
        img_urls: List of URLs to the image files (can be generated using generate_tmp_public_url for local files)
        user_input: Question or instruction about the image content(s)

    Returns:
        str: Claude's analysis of the image(s) based on the query, or Plotly-compatible data structure if visualization is needed
    """
    print("chat_with_imgs============>", img_urls, user_input)
    return analyze_imgs(img_urls, user_input)


@tool
def generate_tmp_public_url(file_path: str) -> str:
    """
    Generate a temporary public URL for a local file. The file will be automatically deleted after 7 days.

    Args:
        file_path: The path to the local file you want to make publicly accessible

    Returns:
        str: A public URL that can be used to access the file for 7 days

    Raises:
        FileNotFoundError: If the specified file does not exist
    """
    print("generate_tmp_public_url============>", file_path)
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"File not found: {file_path}")

    return upload_and_get_tmp_public_url(
        file_path,
        DICT_VAR.get("botrun_flow_lang_url", ""),
        DICT_VAR.get("user_id", ""),
    )


@tool
def web_search(user_input: str):
    """
    Use this to search the web when you need up-to-date information or when your knowledge is insufficient.
    This tool uses Perplexity to perform web searches and provides detailed answers with citations.

    Args:
        user_input: The search query or question you want to find information about.

    Returns:
        str: A detailed answer based on web search results, including citations to source materials
    """
    print("web_search============>", user_input)
    now = datetime.now()
    dates = format_dates(now)
    western_date = dates["western_date"]
    taiwan_date = dates["taiwan_date"]

    final_input = (
        f"現在的西元時間：{western_date}\n現在的民國時間：{taiwan_date}\n\n{user_input}"
    )
    try:
        # 定義一個內部的非同步函數來處理搜尋結果
        async def process_search():
            search_result = ""
            async for event in respond_with_perplexity_search(
                final_input,
                user_prompt_prefix="",
                messages_for_llm=[],
                domain_filter=[],
                stream=False,
            ):
                if event and isinstance(event.chunk, str):
                    search_result += event.chunk
            return search_result

        # 使用 asyncio.run 執行非同步函數
        search_result = asyncio.run(process_search())
    except Exception as e:
        import traceback

        traceback.print_exc()
        return f"Error during web search: {str(e)}"
    return search_result if search_result else "No results found."


@tool
def generate_image(user_input: str):
    """
    Use this to generate high-quality images using DALL-E 3.

    Capabilities:
    - Creates photorealistic images and art
    - Handles complex scenes and compositions
    - Maintains consistent styles
    - Follows detailed prompts with high accuracy
    - Supports various artistic styles and mediums

    Best practices for prompts:
    - Be specific about style, mood, lighting, and composition
    - Include details about perspective and setting
    - Specify artistic medium if desired (e.g., "oil painting", "digital art")
    - Mention color schemes or specific colors
    - Describe the atmosphere or emotion you want to convey

    Limitations:
    - Cannot generate images of public figures or celebrities
    - Avoids harmful, violent, or adult content
    - May have inconsistencies with hands, faces, or text
    - Cannot generate exact copies of existing artworks or brands
    - Limited to single image generation per request

    Args:
        user_input: Detailed description of the image you want to generate.
                   Be specific about style, content, and composition.

    Returns:
        str: URL to the generated image, or error message if generation fails
    """
    try:
        print("generate_image============>", user_input)
        image_response = image_generation(
            prompt=user_input,
            model="dall-e-3",
            api_key=os.getenv("OPENAI_API_KEY"),
        )

        image_url = image_response["data"][0]["url"]
        print("generate_image generated============>", image_url)
        return image_url
    except Exception as e:
        return f"Error: {e}"


@tool
def create_plotly_chart(figure_data: dict, title: str = None) -> str:
    """
    Create an interactive Plotly visualization and return its URL.
    This URL should be provided to the user,
    as they will need to access it to view the interactive chart in their web browser.

    使用 create_plotly_chart 的情境：
    - 提到「圖表」、「統計圖」、「視覺化」等字眼
    - 需要呈現數據趨勢（折線圖）
    - 需要比較數值（長條圖、圓餅圖）
    - 需要展示分布情況（散點圖、熱力圖）
    - 需要顯示時間序列資料（時間軸圖）
    - 需要展示地理資訊（地圖）
    - 需要多維度資料分析（3D圖、氣泡圖）
    - 需要展示統計分布（箱型圖）
    - 需要展示累積趨勢（面積圖）
    - 需要互動式資料探索

    Integration with Other Tools:
    This function can be used in conjunction with chat_with_imgs and chat_with_pdf when they return data
    suitable for visualization. When those tools detect a need for visualization, they will return a dictionary
    with a "__plotly_data__" key, which can be directly passed to this function.

    Example workflow:
    1. User asks to analyze and visualize data from images/PDFs
    2. chat_with_imgs or chat_with_pdf returns data with "__plotly_data__" key
    3. Pass that data to this function to get an interactive visualization URL

    Supported Chart Types:
    - Line charts: For showing trends and time series data
    - Bar charts: For comparing values across categories
    - Pie charts: For showing proportions of a whole
    - Scatter plots: For showing relationships between variables
    - Heat maps: For showing patterns in matrix data
    - Box plots: For showing statistical distributions
    - Geographic maps: For showing spatial data
    - 3D plots: For showing three-dimensional data
    - Bubble charts: For showing three variables in 2D
    - Area charts: For showing cumulative totals over time

    The figure_data should be a dictionary containing plotly figure specifications with 'data' and 'layout'.
    Example:
    {
        'data': [{
            'type': 'scatter',
            'x': [1, 2, 3, 4],
            'y': [10, 15, 13, 17]
        }],
        'layout': {
            'title': 'My Plot'
        }
    }

    Args:
        figure_data: Dictionary containing plotly figure specifications or output from chat_with_imgs/chat_with_pdf
        title: Optional title for the plot. If provided, must be in Traditional Chinese.
               For example: "台灣人口統計圖表" instead of "Taiwan Population Chart"

    Returns:
        str: URL for the interactive HTML visualization. This URL should be provided to the user,
             as they will need to access it to view the interactive chart in their web browser.
    """
    try:
        print("create_plotly_chart============>", figure_data, title)
        # If the input is from chat_with_imgs or chat_with_pdf, extract the plotly data
        if isinstance(figure_data, dict) and "__plotly_data__" in figure_data:
            figure_data = figure_data["__plotly_data__"]

        html_url = generate_plotly_files(
            figure_data,
            DICT_VAR.get("botrun_flow_lang_url", ""),
            DICT_VAR.get("user_id", ""),
            title,
        )
        print("create_plotly_chart generated============>", html_url)
        return html_url
    except Exception as e:
        return f"Error creating visualization URL: {str(e)}"


@tool
def create_mermaid_diagram(mermaid_data: str, title: str = None) -> str:
    """
    Create an interactive Mermaid diagram visualization and return its URL.
    This URL should be provided to the user,
    as they will need to access it to view the interactive diagram in their web browser.

    使用 create_mermaid_diagram 的情境：
    - 提到「流程圖」、「架構圖」、「關係圖」等字眼
    - 需要展示系統架構（flowchart）
    - 需要說明操作流程（flowchart）
    - 需要展示時序互動（sequence diagram）
    - 需要展示狀態轉換（state diagram）
    - 需要展示類別關係（class diagram）
    - 需要展示實體關係（ER diagram）
    - 需要展示專案時程（gantt chart）
    - 需要展示使用者旅程（journey）
    - 需要展示需求關係（requirement diagram）
    - 需要展示資源分配（pie chart）

    Supported Diagram Types:
    1. Flowcharts (graph TD/LR):
       - System architectures
       - Process flows
       - Decision trees
       - Data flows

    2. Sequence Diagrams (sequenceDiagram):
       - API interactions
       - System communications
       - User interactions
       - Message flows

    3. Class Diagrams (classDiagram):
       - Software architecture
       - Object relationships
       - System components
       - Code structure

    4. State Diagrams (stateDiagram-v2):
       - System states
       - Workflow states
       - Process states
       - State transitions

    5. Entity Relationship Diagrams (erDiagram):
       - Database schemas
       - Data relationships
       - System entities
       - Data models

    6. User Journey Diagrams (journey):
       - User experiences
       - Customer flows
       - Process steps
       - Task sequences

    7. Gantt Charts (gantt):
       - Project timelines
       - Task schedules
       - Resource allocation
       - Milestone tracking

    8. Pie Charts (pie):
       - Data distribution
       - Resource allocation
       - Market share
       - Component breakdown

    9. Requirement Diagrams (requirementDiagram):
       - System requirements
       - Dependencies
       - Specifications
       - Constraints

    Example Mermaid syntax for a simple flowchart:
    ```
    graph TD
        A[開始] --> B{是否有資料?}
        B -->|是| C[處理資料]
        B -->|否| D[取得資料]
        C --> E[結束]
        D --> B
    ```

    Args:
        mermaid_data: String containing the Mermaid diagram definition
        title: Optional title for the diagram. If provided, must be in Traditional Chinese.
               For example: "系統流程圖" instead of "System Flowchart"

    Returns:
        str: URL for the interactive HTML visualization. This URL should be provided to the user,
             as they will need to access it to view the interactive diagram in their web browser.
    """
    try:
        print("create_mermaid_diagram============>", mermaid_data, title)
        html_url = generate_mermaid_files(
            mermaid_data,
            DICT_VAR.get("botrun_flow_lang_url", ""),
            DICT_VAR.get("user_id", ""),
            title,
        )
        print("create_mermaid_diagram generated============>", html_url)
        return html_url
    except Exception as e:
        return f"Error creating diagram URL: {str(e)}"


BASIC_TOOLS = [
    current_time,
    scrape,
    days_between,
    chat_with_pdf,
    chat_with_imgs,
    web_search,
    generate_image,
]
DICT_VAR = {}

# Define the graph
from langgraph.prebuilt import create_react_agent

now = datetime.now()
dates = format_dates(now)
western_date = dates["western_date"]
taiwan_date = dates["taiwan_date"]


def create_react_agent_graph(
    system_prompt: str = "",
    botrun_flow_lang_url: str = "",
    user_id: str = "",
):
    """
    Create a react agent graph with optional system prompt

    Args:
        system_prompt: The system prompt to use for the agent
    """
    tools = BASIC_TOOLS
    if botrun_flow_lang_url and user_id:
        DICT_VAR["botrun_flow_lang_url"] = botrun_flow_lang_url
        DICT_VAR["user_id"] = user_id
        tools.append(generate_tmp_public_url)
        tools.append(create_plotly_chart)
        tools.append(create_mermaid_diagram)
        print("tools============>", tools)
    new_system_prompt = (
        system_prompt
        + """

    以下為回應時需注意的準則，請你務必遵守，你會遵守<使用工具 (tools) 需注意的事項>、以及<回應時需注意的準則>。
<回應時需注意的準則>
- 如果你要傳給使用者任何的 URL，請確保 URL 是從 tool 回傳的，或是從歷史記錄中找到的，千萬不可自己創造 URL 回傳給使用者。
</<回應時需注意的準則>>
<使用工具 (tools) 需注意的事項>
- generate_image: 
    當使用 generate_image 工具時，你必須在回應中包含圖片網址。
    請按照以下格式回應(從 @begin img開始，到 @end 結束，中間包含圖片網址)：
    @begin img("{image_url}") @end

- web_search:
    這個工具會回傳參考來源，也就是它回應的內容，參考哪些網站的內容，當你回覆內容給使用者時，也要把相關參考來源回傳給使用者。
    它會回傳類似這樣子的內容：
    {搜尋結果的文字內容}
    參考來源：
    - [{source_title_1}]({source_url_1})
    - [{source_title_2}]({source_url_2})
    - [{source_title_3}]({source_url_3})
    
    你回覆使用者的時候，也請按照以下格式回應：
    {你要回應給使用者的文字內容}
    參考來源：
    - [{source_title_1}]({source_url_1})
    - [{source_title_2}]({source_url_2})
    - [{source_title_3}]({source_url_3})

- chat_with_pdf:
    使用者的問題或指令。如果問題中包含以下關鍵字，要在 user_input 中包含以下字句：
    - 「圖表」、「統計圖」、「視覺化」：幫我生成適合的統計圖表，並提供相關 plotly 的資料給我
    - 「趨勢」、「走向」：幫我生成折線圖，並提供相關 plotly 的資料給我
    - 「比較」、「對照」：幫我生成長條圖或圓餅圖，並提供相關 plotly 的資料給我
    - 「分布」、「分散」：幫我生成散點圖或熱力圖，並提供相關 plotly 的資料給我
    - 「時間序列」：幫我生成時間軸圖，並提供相關 plotly 的資料給我
    - 「地理資訊」：幫我生成地圖，並提供相關 plotly 的資料給我
    - 「多維度分析」：幫我生成 3D 圖或氣泡圖，並提供相關 plotly 的資料給我
    - 「流程圖」、「流程」：幫我生成 flowchart，並提供相關 mermaid 的資料給我
    - 「架構圖」、「架構」：幫我生成 flowchart，並提供相關 mermaid 的資料給我
    - 「關係圖」、「關係」：幫我生成 flowchart 或 ER diagram，並提供相關 mermaid 的資料給我
    - 「時序圖」、「序列圖」：幫我生成 sequence diagram，並提供相關 mermaid 的資料給我
    - 「狀態圖」、「狀態」：幫我生成 state diagram，並提供相關 mermaid 的資料給我
    - 「類別圖」、「類別」：幫我生成 class diagram，並提供相關 mermaid 的資料給我
    - 「甘特圖」、「時程圖」：幫我生成 gantt chart，並提供相關 mermaid 的資料給我

- chat_with_imgs:
    使用者的問題或指令。如果問題中包含以下關鍵字，要在 user_input 中包含以下字句：
    - 「圖表」、「統計圖」、「視覺化」：幫我生成適合的統計圖表，並提供相關 plotly 的資料給我
    - 「趨勢」、「走向」：幫我生成折線圖，並提供相關 plotly 的資料給我
    - 「比較」、「對照」：幫我生成長條圖或圓餅圖，並提供相關 plotly 的資料給我
    - 「分布」、「分散」：幫我生成散點圖或熱力圖，並提供相關 plotly 的資料給我
    - 「時間序列」：幫我生成時間軸圖，並提供相關 plotly 的資料給我
    - 「地理資訊」：幫我生成地圖，並提供相關 plotly 的資料給我
    - 「多維度分析」：幫我生成 3D 圖或氣泡圖，並提供相關 plotly 的資料給我
    - 「流程圖」、「流程」：幫我生成 flowchart，並提供相關 mermaid 的資料給我
    - 「架構圖」、「架構」：幫我生成 flowchart，並提供相關 mermaid 的資料給我
    - 「關係圖」、「關係」：幫我生成 flowchart 或 ER diagram，並提供相關 mermaid 的資料給我
    - 「時序圖」、「序列圖」：幫我生成 sequence diagram，並提供相關 mermaid 的資料給我
    - 「狀態圖」、「狀態」：幫我生成 state diagram，並提供相關 mermaid 的資料給我
    - 「類別圖」、「類別」：幫我生成 class diagram，並提供相關 mermaid 的資料給我
    - 「甘特圖」、「時程圖」：幫我生成 gantt chart，並提供相關 mermaid 的資料給我

- create_plotly_chart:
    當使用 create_plotly_chart 工具時，你必須在回應中包含create_plotly_chart回傳的URL網址。
    請按照以下格式回應：
    [{plotly_chart_title}] ({plotly_chart_url})
    <範例1>
    使用者提問：
    請幫我分析這個PDF的內容，並產出一個圖表給我看
    回應：
    {分析的內容文字}
    我為你製作了一個圖表，請看這個網址：
    [{plotly_chart_title}] ({plotly_chart_url})
    </範例1>
    <範例2>
    使用者提問：
    請幫我深度分析這個檔案內容，並產出讓使用者好懂的比較圖
    回應：
    {分析的內容文字}
    我為你製作了一個圖表，請看這個網址：
    [{plotly_chart_title}] ({plotly_chart_url})
    </範例2>

- create_mermaid_diagram:
    當使用 create_mermaid_diagram 工具時，你必須在回應中包含create_mermaid_diagram回傳的URL網址。
    請按照以下格式回應：
    [{mermaid_diagram_title}] ({mermaid_diagram_url})
    <範例1>
    使用者提問：
    請幫我分析這個PDF的內容，根據開會決議，產出一個行動流程圖
    回應：
    {分析的內容文字}
    我為你製作了一個流程圖，請看這個網址：
    [{mermaid_diagram_title}] ({mermaid_diagram_url})
    </範例1>
    <範例2>
    使用者提問：
    請幫我深度分析這個檔案內容，根據新的伺服器架構，產出一個架構圖
    回應：
    {分析的內容文字}
    我為你製作了一個架構圖，請看這個網址：
    [{mermaid_diagram_title}] ({mermaid_diagram_url})
    </範例2>
</使用工具 (tools) 需注意的事項>
    """
    )
    return create_react_agent(
        model,
        tools=tools,
        prompt=new_system_prompt,
        checkpointer=MemorySaver(),
    )


# Default graph instance with empty prompt
graph = create_react_agent_graph()
# LangGraph Studio 測試用
# graph = create_react_agent_graph(
#     system_prompt="",
#     botrun_flow_lang_url="https://botrun-flow-lang-fastapi-dev-36186877499.asia-east1.run.app",
#     user_id="sebastian.hsu@gmail.com",
# )
