from pydantic import Field
from utils import _err, _log, _log_with_details, load_object, store_object

from melvin.tecton_gen_ai.api import Agent, AgentInputModel


class Input(AgentInputModel):
    query: str = Field(
        description="The detailed instruction on how to generate the chart"
    )
    input_session_object_key: str = Field(
        description="The session object key of input data from calling a dataframe generator, the key must start with `so_df_`",
    )


def sys_prompt(input_session_object_key: str) -> str:
    import datetime

    try:
        description = load_object(input_session_object_key)["description"]
    except Exception as e:
        return (
            f"There is an error with the input data key: {input_session_object_key}. Error: {e} "
            "Do not call any tool, directly return this error message"
        )
    time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    return f"""
You are going to generate a single chart based on user's instruction and the input data.

You need to first transform the data using duckdb SQL by `run_sql`.

The data transformation should take the chart type into consideration

After the transformation, call one and only one chart tool based on the user's intent to output the chart image.

If not successful, you should return an error message starting with `Error: `
If successful, return a chart representation which is a session object key starting with `so_chart_`,
and surround it with triple backticks, language is `tecton_chart`

The final chart must be generated by one of the chart tools. don't manipulate the `so_chart_` id.

Example of the chart representation:

```tecton_chart
so_chart_sdgg
```

Note:

    - strptime is for string to time conversion
    - strftime is for time to string conversion
    - a period of time from now (e.g. since 3 days ago) can be expressed as: `current_localtimestamp() - INTERVAL 3 DAY`

When the chart type is not explicitly specified in the query, you should choose the chart type based on these guidelines:

    - Line chart: for timeseries, and each line represents a column of the data
    - Line chart with categories: for timeseries, and each line represents a value in a category column of the data
    - Stacked area chart: for timeseries, similar as line chart with categories but when a breakdown is specified
    - Pie chart: not for timeseries, but a breakdown on values of a column

Current time:

{time}

The generated duckdb SQL should only use the columns in the description and schema of the input data.
Here are the description and schema of the input data:

{description}
"""


def run_sql(input_session_object_key: str, duckdb_sql: str) -> str:
    """
    Run duckdb SQL to transform the data before calling a chart tool.
    It should assume the query is operating on a table with name: `input_data`
    And the description and schema of `input_data` is the same as the input data

    Args:

        input_session_object_key: the key to the input data
        duckdb_sql: the duckdb_sql to run, the sql must be a SELECT statement

    Returns:

        str: the session object key of the transformed data, starting with `so_charttemp_` or an error message
    """
    import duckdb

    _log_with_details(
        ":question: Transforming data using duckdb", f"```sql\n{duckdb_sql}"
    )
    try:
        input_data = load_object(input_session_object_key)["df"]  # noqa
        input_row = len(input_data)
        with duckdb.connect() as con:
            output_df = con.sql(duckdb_sql).df()
            output_row = len(output_df)
        _log(
            "Data transformed. "
            f"Input data has {input_row} rows, output data has {output_row} rows"
        )
        return store_object(output_df, "charttemp")
    except Exception as e:
        return _err(e)


def generate_line_chart(
    transformed_data_key: str, x_label: str, y_labels: list[str], title: str
) -> str:
    """
    Generate line chart. `run_sql` must be called to transform the original data

    Args:

        transformed_data_key: the session object key generated by `run_sql`, starting with `so_charttemp_`
        x_label: label of x axis
        y_labels: list of labels of the lines
        title: the title of the chart

    Returns:

        str: the session object key of the output image, starting with `so_chart_`

    Note:

        - The first column should be the x axis
        - The rest columns will be on the y axis
    """
    import plotly.express as px
    import plotly.io as pio

    _log(f":chart: Generating line chart: {title}")

    df = load_object(transformed_data_key)
    df.columns = [x_label] + y_labels
    fig = px.line(
        df,
        x=df.columns[0],
        y=df.columns[1:],
        title=title,
    )
    return store_object(pio.to_json(fig), "chart")


def generate_line_chart_with_categories(
    transformed_data_key: str,
    x_label: str,
    y_label: str,
    category_label: str,
    title: str,
) -> str:
    """
    Generate line chart with the category where the categories are also from the dataframe.
    `run_sql` must be called to transform the original data

    Args:

        transformed_data_key: the session object key generated by `run_sql`, starting with `so_charttemp_`
        x_label: label of x axis
        y_label: label of y axis
        category_label: label of category axis
        title: the title of the chart

    Returns:

        str: the session object key of the output image, starting with `so_chart_`

    Note on the data transformation:

        - The first column should be the x axis
        - The last column should be the value given each x and categories
        - The middle columns should be the category columns

    """
    import pandas as pd
    import plotly.express as px
    import plotly.io as pio

    _log(f":chart: Generating line chart with categories: {title}")
    df = load_object(transformed_data_key)
    if len(df.columns) < 3:
        return _err("There should be at least 3 columns for this chart")
    elif len(df.columns) == 3:
        df.columns = [x_label, category_label, y_label]
    else:
        str_col = df[df.columns[1]].astype(str)
        for i in range(2, len(df.columns) - 1):
            str_col = str_col.str.cat(df[df.columns[i]].astype(str), sep=" - ")
        df = pd.DataFrame(
            {
                x_label: df[df.columns[0]],
                category_label: str_col,
                y_label: df[df.columns[-1]],
            }
        )
    fig = px.line(
        df,
        x=df.columns[0],
        y=df.columns[-1],
        color=df.columns[1],
        title=title,
    )
    return store_object(pio.to_json(fig), "chart")


def generate_stacked_area_chart(
    transformed_data_key: str,
    x_label: str,
    y_label: str,
    category_label: str,
    title: str,
) -> str:
    """
    Generate stacked area chart with the category where the categories are also from the dataframe.
    `run_sql` must be called to transform the original data

    Args:

        transformed_data_key: the session object key generated by `run_sql`, starting with `so_charttemp_`
        x_label: label of x axis
        y_label: label of y axis
        category_label: label of category axis
        title: the title of the chart

    Returns:

        str: the session object key of the output image, starting with `so_chart_`

    Note on the data transformation:

        - The first column should be the x axis
        - The last column should be the value given each x and categories
        - The middle columns should be the category columns

    """
    import pandas as pd
    import plotly.express as px
    import plotly.io as pio

    _log(f":chart: Generating stacked area chart: {title}")
    df = load_object(transformed_data_key)
    if len(df.columns) < 3:
        return _err("There should be at least 3 columns for this chart")
    elif len(df.columns) == 3:
        df.columns = [x_label, category_label, y_label]
    else:
        str_col = df[df.columns[1]].astype(str)
        for i in range(2, len(df.columns) - 1):
            str_col = str_col.str.cat(df[df.columns[i]].astype(str), sep=" - ")
        df = pd.DataFrame(
            {
                x_label: df[df.columns[0]],
                category_label: str_col,
                y_label: df[df.columns[-1]],
            }
        )
    fig = px.area(
        df,
        x=df.columns[0],
        y=df.columns[-1],
        color=df.columns[1],
        title=title,
    )
    return store_object(pio.to_json(fig), "chart")


def generate_pie_chart(transformed_data_key: str, title: str) -> str:
    """
    Generate pie chart. `run_sql` must be called to transform the original data

    Args:

        transformed_data_key: the session object key generated by `run_sql`, starting with `so_charttemp_`
        title: the title of the chart

    Returns:

        str: the session object key of the output image, starting with `so_chart_`

    Note:

        - The first column should be the labels
        - The second column should be the values
    """
    import plotly.express as px
    import plotly.io as pio

    _log(f":chart: Generating pie chart: {title}")

    df = load_object(transformed_data_key)
    if len(df.columns) != 2:
        return _err("There should be exactly 2 columns for this chart")
    fig = px.pie(df, names=df.columns[0], values=df.columns[1], title=title)
    return store_object(pio.to_json(fig), "chart")


def _make_chart_agent() -> Agent:
    return Agent(
        name="chart_agent",
        description="""An agent that generates charts given instruction and input data

Returns:

    if not successful, it returns an error message starting with `Error: `
    if successful, it returns the chart type (e.g. line chart) and a chart representation quoted with triple backticks, language is `tecton_chart`

Note:

    - A dataframe generator tool must be called before calling this agent
    - The output of this agent should never be used to call this agent again
    - For any follow up question, the input data needs to be regenerated
    - Do not try to interpret the content in a tecton_chart code block, use the exact code block in the final output.
""",
        prompt=sys_prompt,
        input_schema=Input,
        tools=[
            run_sql,
            generate_line_chart,
            generate_line_chart_with_categories,
            generate_pie_chart,
            generate_stacked_area_chart,
        ],
    )
