import os
import json
from vanna.ollama import Ollama
from vanna.qianwen import QianWenAI_Chat
from vanna.deepseek import DeepSeekChat
from extensions.utils.rewrite_ask import ask
from dotenv import load_dotenv
import plotly.io as pio
from vanna.milvus import Milvus_VectorStore
from pymilvus import MilvusClient,model
from collections import defaultdict


load_dotenv()
# 设置显示后端为浏览器
pio.renderers.default = 'browser'
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from typing import List
import ollama
import numpy as np
from pymilvus.model.base import BaseEmbeddingFunction
# 自定义嵌入式模型（适配milvus向量数据库）
class CustomEmbeddingFunction(BaseEmbeddingFunction):
    def __init__(self, config=None):
        model_host = config['host'] if "host" in config else 'http://wsd.wisdomidata.com:19042'
        self.embed_model = config['embed_model'] if "embed_model" in config else 'bge-m3'
        self.embedding_model = ollama.Client(model_host)
        self.keep_alive = config.get('keep_alive', None)
        self.ollama_options = config.get('options', {})
        self.num_ctx = self.ollama_options.get('num_ctx', 2048)

    def __call__(self, texts: List[str]):
        self._encode(texts)
    def _encode(self,texts: list[str]) -> list[list[float]]:
        return [self.embedding_model.embeddings(
            model=self.embed_model,
            prompt=text,
            options=self.ollama_options,
            keep_alive=self.keep_alive
        )["embedding"] for text in texts]
    def encode_documents(self, documents: List[str]) -> List[np.array]:
        # 将每个嵌入结果转换为 np.ndarray
        embeddings = self._encode(documents)
        return [np.array(embedding) for embedding in embeddings]
    def encode_queries(self, queries: List[str]) -> List[np.array]:
        embeddings = self._encode(queries)
        return [np.array(embedding) for embedding in embeddings]

class VannaServer:
    def __init__(self, config):
        self.config = config
        self.vn = self._initialize_vn()

    def _initialize_vn(self):
        config = self.config
        supplier = config["supplier"]
        llm_type = config["llm_type"]
        model_ = config["model"]
        api_key = config["api_key"]
        ollama_host = config["ollama_host"] if "ollama_host" in config else None
        milvus_uri = config["milvus_uri"]
        sql_type = config["sql_type"]
        host = config["host"] if "host" in config else os.getenv("DB_HOST", "localhost")
        dbname = config["dbname"] if "dbname" in config else os.getenv("DB_NAME", "dify_data")
        user = config["user"] if "user" in config else os.getenv("DB_USER", "root")
        password = config["password"] if "password" in config else os.getenv("DB_PASSWORD", "mysql")
        port = config["port"] if "port" in config else int(os.getenv("DB_PORT", 3306))
        milvus_database = config["milvus_database"] if "milvus_database" in config else "test"
        milvus_client = MilvusClient(uri=milvus_uri,db_name=milvus_database)

        embedding_host = config["embedding_host"] if "embedding_host" in config else 'http://wsd.wisdomidata.com:19042'
        embedding_model = config["embedding_model"] if "embedding_model" in config else "bge-m3" # BAAI/bge-m3
        embedding_function = CustomEmbeddingFunction({
            "host": embedding_host,
            "embed_model": embedding_model
        })
        chat_llm = Ollama
        if llm_type == "ollama":
            config = {
                'model': model_,  # 本地ollama大模型名称
                'ollama_host': ollama_host,  # 本地ollama大模型服务地址
                'milvus_client': milvus_client,  # 本地milvus向量数据库服务地址
                "n_results": 12,
                "embedding_function": embedding_function,
            }
        else:
            config = {
                'model': model_,  # 本地ollama大模型名称
                'api_key': api_key,  # 本地ollama大模型服务地址
                'milvus_client': milvus_client,  # 本地milvus向量数据库服务地址
                "n_results": 12,
                "embedding_function": embedding_function,
            }
        if llm_type == "tongyi":
            chat_llm = QianWenAI_Chat
        elif llm_type == "deepseek":
            chat_llm = DeepSeekChat

        MyVanna = make_vanna_class(ChatClass=chat_llm)
        vn = MyVanna(config)
        if sql_type == "postgres":
            vn.connect_to_postgres(host=host, dbname=dbname, user=user, password=password, port=port)
        elif sql_type == "mysql":
            vn.connect_to_mysql(host=host, dbname=dbname, user=user, password=password, port=port)

        return vn

    def schema_train(self):
        # The information schema query may need some tweaking depending on your database. This is a good starting point.
        df_information_schema = self.vn.run_sql("SELECT * FROM INFORMATION_SCHEMA.COLUMNS where table_schema = 'public'")

        # This will break up the information schema into bite-sized chunks that can be referenced by the LLM
        plan = self.vn.get_training_plan_generic(df_information_schema)
        # print(plan)

        # If you like the plan, then uncomment this and run it to train
        self.vn.train(plan=plan)

    # 更新建表DDL语句
    def refresh_create_table_ddl_train(self):
        sql = """
SELECT 
    'CREATE TABLE ' 
    || C.TABLE_NAME 
    || ' (' 
    || C.COLUMN_NAMES 
    || ');'
    || C.COMMENT_COLUMNS
		|| CASE WHEN FK.FOREIGN_KEY_COLUMNS IS NOT NULL THEN FK.FOREIGN_KEY_COLUMNS ELSE '' END
		|| CASE WHEN FK.FOREIGN_KEY_DESC IS NOT NULL THEN FK.FOREIGN_KEY_DESC ELSE '' END
    || 'COMMENT ON TABLE ' 
    || C.TABLE_NAME 
    || ' IS ''' 
    || G.DESCRIPTION 
    || ''';'
		AS DDL,
    C.TABLE_NAME
FROM (
    SELECT 
        COL.TABLE_NAME,
        COL.TABLE_SCHEMA,
        STRING_AGG(
            COL.COLUMN_NAME 
            || ' ' 
            || COL.DATA_TYPE 
            || COALESCE('(' || COL.CHARACTER_MAXIMUM_LENGTH || ')', '') 
            || COALESCE(' DEFAULT ' || COL.COLUMN_DEFAULT, '') 
            || CASE 
                WHEN COL.IS_NULLABLE = 'NO' THEN ' NOT NULL' 
                ELSE '' 
              END,
            ','
        ) AS COLUMN_NAMES,
        STRING_AGG(
            'COMMENT ON COLUMN ' 
            || COL.TABLE_NAME 
            || '.' 
            || COL.COLUMN_NAME 
            || ' IS ''' 
            || PGD.DESCRIPTION 
            || ''';',
            ''
        ) AS COMMENT_COLUMNS
    FROM 
        PG_CATALOG.PG_STATIO_ALL_TABLES AS ST
    INNER JOIN 
        PG_CATALOG.PG_DESCRIPTION AS PGD 
        ON PGD.OBJOID = ST.RELID
    INNER JOIN 
        INFORMATION_SCHEMA.COLUMNS AS COL 
        ON (
            COL.TABLE_SCHEMA = ST.SCHEMANAME 
            AND COL.TABLE_NAME = ST.RELNAME 
            AND COL.ORDINAL_POSITION = PGD.OBJSUBID
        )
    WHERE 
        COL.TABLE_SCHEMA = 'public'
    GROUP BY 
        COL.TABLE_SCHEMA, 
        COL.TABLE_NAME
) C
LEFT JOIN (
    SELECT 
        N.NSPNAME AS SCHEMA_NAME,
        C.RELNAME AS TABLE_NAME,
        D.DESCRIPTION
    FROM 
        PG_CATALOG.PG_DESCRIPTION D
    JOIN 
        PG_CATALOG.PG_CLASS C 
        ON C.OID = D.OBJOID
    JOIN 
        PG_CATALOG.PG_NAMESPACE N 
        ON N.OID = C.RELNAMESPACE
    WHERE 
        C.RELKIND = 'r'
        AND D.OBJSUBID = 0
) G 
ON G.SCHEMA_NAME = C.TABLE_SCHEMA 
AND G.TABLE_NAME = C.TABLE_NAME
LEFT JOIN (
    SELECT rel_src.relname AS source_table,
        STRING_AGG(
            'ALTER TABLE ' 
            || rel_src.relname
            || ' ADD CONSTRAINT ' 
            || con.conname 
            || ' FOREIGN KEY ('
            || att_src.attname 
            || ') REFERENCES '
            || rel_tgt.relname
            || '('
            || att_tgt.attname
            || ');'
            ,
            ''
      ) AS FOREIGN_KEY_COLUMNS,
        STRING_AGG(
                'COMMENT ON CONSTRAINT  ' 
                || con.conname 
                || ' ON ' 
                || rel_src.relname 
                || ' IS ''' 
                || d.description
                || ''';',
                ''
        ) AS FOREIGN_KEY_DESC
    FROM
        pg_constraint con
        JOIN pg_class rel_src ON rel_src.oid = con.conrelid
        JOIN pg_class rel_tgt ON rel_tgt.oid = con.confrelid
        JOIN pg_attribute att_src ON att_src.attrelid = rel_src.oid AND att_src.attnum = ANY(con.conkey)
        JOIN pg_attribute att_tgt ON att_tgt.attrelid = rel_tgt.oid AND att_tgt.attnum = ANY(con.confkey)
        LEFT JOIN pg_description d ON d.objoid = con.oid
    WHERE
        con.contype = 'f'
    GROUP BY
        rel_src.relname
) FK ON FK.source_table = C.TABLE_NAME
WHERE C.TABLE_NAME NOT IN ('flyway_table_dict','flyway_schema_history')
"""
        # The information schema query may need some tweaking depending on your database. This is a good starting point.
        c_table_ddl_list = self.vn.run_sql(sql)

        # 将 DataFrame 转换为字典列表
        c_table_ddl_records = c_table_ddl_list.to_dict(orient='records')

        exist_ddl_data = self.vn.milvus_client.query(
            collection_name="vannaddl",
            output_fields=["*"],
            limit=10000,
        )
        exists_list = filter(lambda m: m["ddl"].startswith("CREATE TABLE "), exist_ddl_data)
        remove_ids = [exist["id"] for exist in exists_list]
        if len(remove_ids) > 0:
            self.vn.milvus_client.delete(collection_name="vannaddl", ids=remove_ids)

        for table_ddl in c_table_ddl_records:
            self.vn.train(ddl=table_ddl["ddl"])

        self.vn.milvus_client.refresh_load(collection_name="vannaddl")


    def refresh_schema_train(self):
        exist_doc_data = self.vn.milvus_client.query(
            collection_name="vannadoc",
            output_fields=["*"],
            limit=10000,
        )
        exists_list = filter(lambda m: m["doc"].startswith("The following columns are in the "), exist_doc_data)
        remove_ids = [exist["id"] for exist in exists_list]
        if len(remove_ids) > 0:
            self.vn.milvus_client.delete(collection_name="vannadoc", ids=remove_ids)
        self.schema_train()
        self.vn.milvus_client.refresh_load(collection_name="vannadoc")

    def update_schema_train_list(self,docs : list[str]):
        exist_doc_data = self.vn.milvus_client.query(
            collection_name="vannadoc",
            output_fields=["*"],
            limit=10000,
        )
        exists_list = filter(lambda m: not m["doc"].startswith("The following columns are in the "), exist_doc_data)
        remove_ids = [exist["id"] for exist in exists_list]
        if len(remove_ids) > 0:
            self.vn.milvus_client.delete(collection_name="vannadoc", ids=remove_ids)
        dict_docs = self.get_dict_docs()
        docs.extend(dict_docs)

        for doc in docs:
            self.vn.train(documentation=doc)
        # self.schema_train()
        self.vn.milvus_client.refresh_load(collection_name="vannadoc")

    def get_dict_docs(self) -> list[str]:
        dict_docs = []
        sql = "select id,table_name,column_name,column_remark,table_remark,dict_values from flyway_table_dict"
        c_table_dict_list = self.vn.run_sql(sql)
        # 将 DataFrame 转换为字典列表
        c_table_dict_records = c_table_dict_list.to_dict(orient='records')

        table_names = list(set(item['table_name'] for item in c_table_dict_records))

        grouped = defaultdict(list)
        for table_dict in c_table_dict_records:
            table_name = table_dict['table_name']  # 分组依据字段
            grouped[table_name].append(table_dict)

        grouped_dict = dict(grouped)

        for table_name in table_names:
            columns_list = grouped_dict[table_name]
            dict_values = ';'.join(f"字段:{item['column_remark']}({item['column_name']})的值:{item["dict_values"]}" for item in columns_list)
            column = columns_list[0]
            doc = f"{column["table_remark"]}表:{column["table_name"]},{dict_values}"
            dict_docs.append(doc)
        return dict_docs

    def vn_train(self, question="", sql="", documentation="", ddl=""):
        if question and sql:
            # 训练问答对
            self.vn.train(
                question=question,
                sql=sql
            )
        elif sql:
            # You can also add SQL queries to your training data. This is useful if you have some queries already laying around. You can just copy and paste those from your editor to begin generating new SQL.
            self.vn.train(sql=sql)

        if documentation:
            # Sometimes you may want to add documentation about your business terminology or definitions.
            self.vn.train(documentation=documentation)

        if ddl:
            # You can also add DDL queries to your training data. This is useful if you have some queries already laying around. You can just copy and paste those from your editor to begin generating new SQL.
            self.vn.train(ddl=ddl)

    def get_training_data(self):
        training_data = self.vn.get_training_data()
        # print(training_data)
        return training_data

    def ask(self, question, visualize=True, auto_train=True, *args, **kwargs):
        sql, df, fig = ask(self.vn, question, visualize=visualize, auto_train=auto_train, *args, **kwargs)
        return sql, df, fig

    def generate_sql(self, question):
        return self.vn.generate_sql(question=question)

    def run_sql(self, sql):
        return self.vn.run_sql(sql=sql)

    def training_data_export(self):
        training_data = self.vn.milvus_client.query(
            collection_name="vannasql",
            output_fields=["*"],
            limit=10000,
        )
        result = []
        if training_data is not None:
            result = [{"question":t['text'], "sql": t['sql']} for t in training_data]

        return result

    def training_data_import(self, data_list):

        empty_items = list(filter(
            lambda item: item['question'] is None or item['question'] == "" or item['sql'] is None or item['sql'] == "",
            data_list
        ))

        if bool(empty_items):
            return True

        exist_doc_data = self.vn.milvus_client.query(
            collection_name="vannasql",
            output_fields=["*"],
            limit=10000,
        )
        data_texts = {t["question"]: t for t in data_list}
        
        if bool(exist_doc_data):
           remove_ids = [item["id"] for item in exist_doc_data if item['text'] in data_texts ]
        
           if bool(remove_ids):
               self.vn.milvus_client.delete(collection_name="vannasql", ids=remove_ids)
        
        for item in data_list:
            self.vn.train(
                question=item["question"],
                sql=item["sql"],
            )
        
        self.vn.milvus_client.refresh_load(collection_name="vannasql")

        return False

def make_vanna_class(ChatClass=Ollama):
    class MyVanna(Milvus_VectorStore, ChatClass):
        def __init__(self, config=None):
            Milvus_VectorStore.__init__(self, config=config)
            ChatClass.__init__(self, config=config)

        def is_sql_valid(self, sql: str) -> bool:
            # Your implementation here
            return False

        def generate_query_explanation(self, sql: str):
            my_prompt = [
                self.system_message("You are a helpful assistant that will explain a SQL query"),
                self.user_message("Explain this SQL query: " + sql),
            ]

            return self.submit_prompt(prompt=my_prompt)

    return MyVanna


# 使用示例
if __name__ == '__main__':
    config = {"supplier": "GITEE"}
    server = VannaServer(config)
    # server.schema_train()
    server.ask("汇总每个类别的销售量和销售额, 并按照销售量进行降序排列")