from pandasai import Agent
from pydantic import BaseModel
import json
from web.file import FileBean, get_dataframes
from base.no_parser import NoParser
from pathlib import Path
from base.init import llm
from web.file import FileBean, get_dataframes
import matplotlib
import matplotlib.pyplot as plt
plt.rcParams['font.family'] = ['Heiti TC', 'Arial Unicode MS']
matplotlib.use('Agg')
config = {"llm": llm, "verbose": True, "open_charts": False,
          "response_parser": NoParser, "save_charts": True}

class BaseReq(BaseModel):
    prompt: str


class PlotReq(BaseReq):
    fileList: list[FileBean]

def gen_plot(req: PlotReq):
    dfs = [df for fb in req.fileList for df in get_dataframes(fb).values()]
    lake = Agent(
        dfs, config={**config, "save_charts_path": f".data/static/plot"})
    response = lake.chat(req.prompt)
    # check response is str
    if isinstance(response, str):
        return {"resultType": 'error', "resultValueStr": response}
    result_type = response['type']
    if result_type == 'dataframe':
        v_dict = response['value'].to_dict(orient='records')
        result_value_str = json.dumps(v_dict)
    elif result_type == 'plot':
        file_path = Path(response["value"])
        result_value_str = file_path.name
    elif result_type == 'number':
        result_value_str = str(response["value"])
    else:
        result_value_str = response["value"]
    return {"resultType": result_type, "resultValueStr": result_value_str}