from typing import List
import plotly.io as pio
import plotly.express as px
import pandas as pd
from utils import TEMP_DIR
import os
import ast
from dotenv import load_dotenv

load_dotenv()

root_url = os.getenv("ROOT_URL", "")


def _write_chart(fig, chart_path, chart_url):
    """Write a Plotly figure to disk and return a responsive iframe HTML string."""
    pio.write_html(fig, chart_path, full_html=False, config={"responsive": True})
    return (
        'Please display this iframe: '
        '<div style="width:100%;overflow-x:auto;">'
        '<iframe style="width:100%;min-width:400px;" height="500" '
        f'src="{chart_url}" frameborder="0" allowfullscreen>'
        '</iframe></div>'
    )


def llm_chart_data_scrub(data, layout):
   #Processing data to account for variation from LLM
   data_list = []
   layout_dict = {}

   if isinstance(data, list):
      data_list = data
   else:
      data_list.append(data) 

   false_replace = [':false', ': false']
   false_value = ':False'   
   true_replace = [':true', ': true']
   true_value = ':True' 

   data_dict = {}   
   for data_obj in data_list:   
      if isinstance(data_obj, str):
         data_obj = data_obj.replace("\n", "")
         for replace in false_replace:
            data_obj = data_obj.replace(replace, false_value)
         for replace in true_replace:
            data_obj = data_obj.replace(replace, true_value)
         print(data_obj)
         data_dict = ast.literal_eval(data_obj)
      else:
         data_dict = data_obj

   if layout and isinstance(layout, list):
      layout_obj = layout[0]
   else:
      layout_obj = layout

   if layout_obj and isinstance(layout_obj, str):
      for replace in false_replace:
         layout_obj = layout_obj.replace(replace, false_value)
      for replace in true_replace:
         layout_obj = layout_obj.replace(replace, true_value)
      print(layout_obj)
      layout_dict = ast.literal_eval(layout_obj)
   else:
      layout_dict = layout_obj

   return data_dict, layout_dict

def scatter_chart_fig(df, x_column: List[str], y_column: str, category: str="", trendline: str="", 
                      trendline_options: List[dict]=[{}], marginal_x: str="", marginal_y: str="",
                      size: str=""):
   
   function_args = {"data_frame":df, "x":x_column, "y":y_column}

   if category:
      function_args["color"] = category
   if trendline:
      function_args["trendline"] = trendline
   if marginal_x:
      function_args["marginal_x"] = marginal_x
   if marginal_y:
      function_args["marginal_y"] = marginal_y
   if size:
      df.loc[df[size] < 0, size] = 0
      function_args["size"] = size   
   if trendline_options:
      trendline_options_dict = {}
      if trendline_options and isinstance(trendline_options, list):
         trendline_options_obj = trendline_options[0]
      else:
         trendline_options_obj = trendline_options

      if trendline_options_obj and isinstance(trendline_options_obj, str):
         trendline_options_dict = ast.literal_eval(trendline_options_obj)
      else:
         trendline_options_dict = trendline_options_obj
      function_args["trendline_options"] = trendline_options_dict

   fig = px.scatter(**function_args)  

   return fig

def scatter_chart_generation_func(x_column: List[str], y_column: str, session_hash, session_folder, data: List[dict]=[{}], layout: List[dict]=[{}], 
                                  category: str="", trendline: str="", trendline_options: List[dict]=[{}], marginal_x: str="", marginal_y: str="",
                                  size: str="", **kwargs):
   try:
      dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
      chart_path = f'{dir_path}/chart.html'
      csv_query_path = f'{dir_path}/query.csv'

      df = pd.read_csv(csv_query_path)

      initial_graph = scatter_chart_fig(df, x_column=x_column, y_column=y_column, 
                              category=category, trendline=trendline, trendline_options=trendline_options,
                              marginal_x=marginal_x, marginal_y=marginal_y, size=size)
      
      fig = initial_graph.to_dict()

      print(data)
      print(layout)

      data_dict,layout_dict = llm_chart_data_scrub(data,layout)

      #Applying stylings and settings generated from LLM
      if layout_dict:
         fig["layout"] = layout_dict

      data_ignore = ["x","y","type"]

      if size:
         data_ignore.append("marker")

      for key, value in data_dict.items():
         if key not in data_ignore:
            for data_item in fig["data"]:
               data_item[key] = value
   
      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
      return {"reply": _write_chart(fig, chart_path, chart_url)}

   except Exception as e:
      print("SCATTER PLOT ERROR")
      print(e)
      reply = f"""There was an error generating the Plotly Scatter Plot from {x_column}, {y_column}, {data}, and {layout}
            The error is {e},
            You should probably try again.
            """
      return {"reply": reply}
   
def line_chart_generation_func(x_column: str, y_column: str, session_hash, session_folder, data: List[dict]=[{}], layout: List[dict]=[{}], 
                                  category: str="", **kwargs):
   try:
      dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
      chart_path = f'{dir_path}/chart.html'
      csv_query_path = f'{dir_path}/query.csv'

      df = pd.read_csv(csv_query_path)

      function_args = {"data_frame":df, "x":x_column, "y":y_column}

      if category:
         function_args["color"] = category

      initial_graph = px.line(**function_args)  

      fig = initial_graph.to_dict()

      data_dict,layout_dict = llm_chart_data_scrub(data,layout)

      print(data_dict)
      print(layout_dict)

      #Applying stylings and settings generated from LLM
      if layout_dict:
         fig["layout"] = layout_dict

      for key, value in data_dict.items():
         if key not in ["x","y","type"]:
            for data_item in fig["data"]:
               data_item[key] = value

      print(fig)

      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
      return {"reply": _write_chart(fig, chart_path, chart_url)}

   except Exception as e:
      print("LINE CHART ERROR")
      print(e)
      reply = f"""There was an error generating the Plotly Line Chart from {x_column}, {y_column}, {data}, and {layout}
            The error is {e},
            You should probably try again.
            """
      return {"reply": reply}

def bar_chart_generation_func(x_column: str, y_column: str, session_hash, session_folder, data: List[dict]=[{}], layout: List[dict]=[{}], 
                                  category: str="", facet_row: str="", facet_col: str="", **kwargs):
   try:
      dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
      chart_path = f'{dir_path}/chart.html'
      csv_query_path = f'{dir_path}/query.csv'

      df = pd.read_csv(csv_query_path)

      function_args = {"data_frame":df, "x":x_column, "y":y_column}

      if category:
         function_args["color"] = category
      if facet_row:
         function_args["facet_row"] = facet_row
      if facet_col:
         function_args["facet_col"] = facet_col

      initial_graph = px.bar(**function_args)  

      fig = initial_graph.to_dict()

      data_dict,layout_dict = llm_chart_data_scrub(data,layout)

      print(data_dict)
      print(layout_dict)

      #Applying stylings and settings generated from LLM
      if layout_dict:
         fig["layout"] = layout_dict

      for key, value in data_dict.items():
         if key not in ["x","y","type"]:
            for data_item in fig["data"]:
               data_item[key] = value

      print(fig)

      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
      return {"reply": _write_chart(fig, chart_path, chart_url)}

   except Exception as e:
      print("BAR CHART ERROR")
      print(e)
      reply = f"""There was an error generating the Plotly Bar Chart from {x_column}, {y_column}, {data}, and {layout}
            The error is {e},
            You should probably try again.
            """
      return {"reply": reply}

def pie_chart_generation_func(values: str, names: str, session_hash, session_folder, data: List[dict]=[{}], layout: List[dict]=[{}], **kwargs):
   try:
      dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
      chart_path = f'{dir_path}/chart.html'
      csv_query_path = f'{dir_path}/query.csv'

      df = pd.read_csv(csv_query_path)

      function_args = {"data_frame":df, "values":values, "names":names}

      initial_graph = px.pie(**function_args)  

      fig = initial_graph.to_dict()

      data_dict,layout_dict = llm_chart_data_scrub(data,layout)

      print(data_dict)
      print(layout_dict)

      #Applying stylings and settings generated from LLM
      if layout_dict:
         fig["layout"] = layout_dict

      for key, value in data_dict.items():
         if key not in ["x","y","type"]:
            for data_item in fig["data"]:
               data_item[key] = value

      print(fig)

      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
      return {"reply": _write_chart(fig, chart_path, chart_url)}

   except Exception as e:
      print("PIE CHART ERROR")
      print(e)
      reply = f"""There was an error generating the Plotly Pie Chart from {values}, {names}, {data}, and {layout}
            The error is {e},
            You should probably try again.
            """
      return {"reply": reply}
   
def histogram_generation_func(x_column: str, session_hash, session_folder, y_column: str="", data: List[dict]=[{}], layout: List[dict]=[{}], histnorm: str="", category: str="",
                              histfunc: str="", **kwargs):   
   try:
      dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
      chart_path = f'{dir_path}/chart.html'
      csv_query_path = f'{dir_path}/query.csv'

      df = pd.read_csv(csv_query_path)

      print(x_column)

      function_args = {"data_frame":df, "x":x_column}

      if y_column:
         function_args["y"] = y_column
      if histnorm:
         function_args["histnorm"] = histnorm
      if category:
         function_args["color"] = category
      if histfunc:
         function_args["histfunc"] = histfunc

      initial_graph = px.histogram(**function_args)  

      fig = initial_graph.to_dict()

      data_dict,layout_dict = llm_chart_data_scrub(data,layout)

      print(data_dict)
      print(layout_dict)

      #Applying stylings and settings generated from LLM
      if layout_dict:
         fig["layout"] = layout_dict

      for key, value in data_dict.items():
         if key not in ["x","y","type"]:
            for data_item in fig["data"]:
               data_item[key] = value

      print(fig)

      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
      return {"reply": _write_chart(fig, chart_path, chart_url)}

   except Exception as e:
      print("HISTOGRAM ERROR")
      print(e)
      reply = f"""There was an error generating the Plotly Histogram from {x_column}.
            The error is {e},
            You should probably try again.
            """
      return {"reply": reply}

def box_chart_generation_func(y_column: str, session_hash, session_folder,
                              x_column: str="", category: str="",
                              layout: List[dict]=[{}], **kwargs):
    try:
        dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
        chart_path = f'{dir_path}/chart.html'
        csv_query_path = f'{dir_path}/query.csv'

        df = pd.read_csv(csv_query_path)

        function_args = {"data_frame": df, "y": y_column}
        if x_column:
            function_args["x"] = x_column
        if category:
            function_args["color"] = category

        initial_graph = px.box(**function_args)
        fig = initial_graph.to_dict()

        _, layout_dict = llm_chart_data_scrub({}, layout)
        if layout_dict:
            fig["layout"] = layout_dict

        chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
        return {"reply": _write_chart(fig, chart_path, chart_url)}

    except Exception as e:
        print("BOX CHART ERROR")
        print(e)
        return {"reply": f"There was an error generating the box plot. Error: {e}. You should probably try again."}


def correlation_heatmap_func(session_hash, session_folder, columns: List[str]=[], **kwargs):
    try:
        dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
        chart_path = f'{dir_path}/chart.html'
        csv_query_path = f'{dir_path}/query.csv'

        df = pd.read_csv(csv_query_path)

        numeric_df = df[columns].select_dtypes(include='number') if columns else df.select_dtypes(include='number')

        if numeric_df.shape[1] < 2:
            return {"reply": "At least two numeric columns are needed for a correlation matrix. Please refine your query to include more numeric columns."}

        corr = numeric_df.corr().round(3)

        fig = px.imshow(
            corr,
            text_auto='.2f',
            color_continuous_scale='RdBu_r',
            zmin=-1,
            zmax=1,
            title='Correlation Matrix',
            aspect='auto',
        )
        fig.update_layout(font=dict(family='Inter, system-ui, sans-serif'))

        chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
        return {"reply": _write_chart(fig, chart_path, chart_url)}

    except Exception as e:
        print("CORRELATION HEATMAP ERROR")
        print(e)
        return {"reply": f"There was an error generating the correlation heatmap. Error: {e}. You should probably try again."}


def rolling_stats_func(x_column: str, y_column: str, session_hash, session_folder,
                       window: int = 7, stats: List[str] = ["mean"],
                       layout: List[dict] = [{}], category: str = "", **kwargs):
    try:
        import plotly.graph_objects as go

        dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
        chart_path = f'{dir_path}/chart.html'
        csv_query_path = f'{dir_path}/query.csv'

        df = pd.read_csv(csv_query_path)

        try:
            df[x_column] = pd.to_datetime(df[x_column])
        except Exception:
            pass
        df = df.sort_values(x_column)

        valid_stats = {"mean", "std", "min", "max"}
        selected_stats = [s for s in stats if s in valid_stats] or ["mean"]

        fig = go.Figure()

        groups = df[category].unique().tolist() if category and category in df.columns else [None]

        for group in groups:
            group_df = df[df[category] == group] if group is not None else df
            prefix = f"{group} — " if group is not None else ""

            fig.add_trace(go.Scatter(
                x=group_df[x_column].values, y=group_df[y_column].values,
                mode="lines", name=f"{prefix}{y_column} (raw)",
                opacity=0.35, line=dict(width=1)
            ))

            rolling_obj = group_df[y_column].rolling(window)
            for stat in selected_stats:
                rolled = getattr(rolling_obj, stat)()
                fig.add_trace(go.Scatter(
                    x=group_df[x_column].values, y=rolled.values,
                    mode="lines", name=f"{prefix}Rolling {stat.capitalize()} (w={window})",
                    line=dict(width=2.5)
                ))

        fig.update_layout(
            title=f"Rolling Statistics (window={window}) — {y_column}",
            xaxis_title=x_column,
            yaxis_title=y_column,
            font=dict(family="Inter, system-ui, sans-serif"),
            legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
        )

        _, layout_dict = llm_chart_data_scrub({}, layout)
        if layout_dict:
            fig.update_layout(**layout_dict)

        chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
        return {"reply": _write_chart(fig, chart_path, chart_url)}

    except Exception as e:
        print("ROLLING STATS ERROR")
        print(e)
        return {"reply": f"There was an error generating the rolling statistics chart. Error: {e}. You should probably try again."}


def table_generation_func(session_hash, session_folder, **kwargs):
    print("TABLE GENERATION")
    try:
        from html import escape

        dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
        csv_query_path = f'{dir_path}/query.csv'

        df = pd.read_csv(csv_query_path)

        total_rows = len(df)
        max_rows = 200
        if total_rows > max_rows:
            df = df.head(max_rows)
            note = (f'<p class="vda-table-note">Showing first {max_rows} of {total_rows} rows'
                    ' — refine your query to see more specific results.</p>')
        else:
            note = ''

        header_cells = ''.join(f'<th>{escape(str(col))}</th>' for col in df.columns)
        row_html = [
            '<tr>' + ''.join(f'<td>{escape(str(val))}</td>' for val in row) + '</tr>'
            for _, row in df.iterrows()
        ]

        style = (
            '<style>'
            '.vda-table-wrap{overflow-x:auto;margin:8px 0;border-radius:8px;border:1px solid #e5e7eb;}'
            '.vda-table{width:100%;border-collapse:collapse;font-size:13px;font-family:Inter,system-ui,sans-serif;}'
            '.vda-table thead th{background:#3B82F6;color:#fff;padding:9px 14px;text-align:left;white-space:nowrap;font-weight:600;}'
            '.vda-table tbody td{padding:7px 14px;border-bottom:1px solid #f1f5f9;white-space:nowrap;}'
            '.vda-table tbody tr:nth-child(even){background:#f8fafc;}'
            '.vda-table tbody tr:last-child td{border-bottom:none;}'
            '.vda-table-note{font-size:12px;color:#6b7280;margin:4px 0 0;text-align:right;}'
            '</style>'
        )

        table = (
            '<div class="vda-table-wrap"><table class="vda-table">'
            f'<thead><tr>{header_cells}</tr></thead>'
            '<tbody>' + '\n'.join(row_html) + '</tbody>'
            '</table></div>'
        )

        return {"reply": style + table + note}

    except Exception as e:
        print("TABLE ERROR")
        print(e)
        return {"reply": f"There was an error generating the table. Error: {e}. You should probably try again."}