import json
from typing import Dict, NamedTuple
import logging
from pilot.out_parser.base import BaseOutputParser, T
from pilot.configs.config import Config
from pilot.scene.chat_db.data_loader import DbDataLoader

CFG = Config()


class SqlAction(NamedTuple):
    sql: str
    thoughts: Dict

    def to_dict(self) -> Dict[str, Dict]:
        return {"sql": self.sql, "thoughts": self.thoughts}


logger = logging.getLogger(__name__)


class DbChatOutputParser(BaseOutputParser):
    def __init__(self, sep: str, is_stream_out: bool):
        super().__init__(sep=sep, is_stream_out=is_stream_out)

    def parse_prompt_response(self, model_out_text):
        clean_str = super().parse_prompt_response(model_out_text)
        print("clean prompt response:", clean_str)
        response = json.loads(clean_str)
        for key in sorted(response):
            if key.strip() == "sql":
                sql = response[key]
            if key.strip() == "thoughts":
                thoughts = response[key]
        return SqlAction(sql, thoughts)

    def parse_view_response(self, speak, data) -> str:
        import pandas as pd

        ### tool out data to table view
        data_loader = DbDataLoader()
        if len(data) < 1:
            data.insert(0, [])
        df = pd.DataFrame(data[1:], columns=data[0])
        if not CFG.NEW_SERVER_MODE and not CFG.SERVER_LIGHT_MODE:
            table_style = """<style> 
                table{border-collapse:collapse;width:100%;height:80%;margin:0 auto;float:center;border: 1px solid #007bff; background-color:#333; color:#fff}th,td{border:1px solid #ddd;padding:3px;text-align:center}th{background-color:#C9C3C7;color: #fff;font-weight: bold;}tr:nth-child(even){background-color:#444}tr:hover{background-color:#444}
             </style>"""
            html_table = df.to_html(index=False, escape=False)
            html = f"<html><head>{table_style}</head><body>{html_table}</body></html>"
            view_text = f"##### {str(speak)}" + "\n" + html.replace("\n", " ")
            return view_text
        else:
            return data_loader.get_table_view_by_conn(data, speak)
