import os

from llama_index.core import SummaryIndex, SQLDatabase, VectorStoreIndex
from llama_index.core.indices.struct_store import SQLTableRetrieverQueryEngine
from llama_index.core.objects import SQLTableNodeMapping, ObjectIndex
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.response_synthesizers import ResponseMode
from sqlalchemy import create_engine

from app.engine import makeDescriptionByEngine
from app.engine.prompt import text_qa_template, refine_template, summary_template, simple_template
from app.engine.retriever.HybridRetriever import HybridRetriever
from app.settings import get_node_postprocessors



def get_Retriever(index,**kwargs):
    bEnableHybrid = True if os.getenv("HYBRID_ENABLED",False).title() == 'True' else False
    if bEnableHybrid:
        alpha = float(os.getenv("HYBRID_ALPHA", "0.5"))
        retriever = HybridRetriever(index,alpha = alpha,**kwargs)
    else:
        retriever = index.as_retriever(**kwargs)
    return retriever


sql_database = None
sql_obj_index = None

# Create a summary query engine
def create_summary_query_engine(top_k=3, use_reranker=False, filters=None):
    global sql_obj_index
    global sql_database
    if sql_obj_index is None or sql_database is None:
        sqlengine = create_engine(os.getenv("SQL_DATABASE_URL", ""))
        sql_database = SQLDatabase(sqlengine)
        table_schema_objs = makeDescriptionByEngine(sql_database)
        table_node_mapping = SQLTableNodeMapping(sql_database)

        sql_obj_index = ObjectIndex.from_objects(
            table_schema_objs,
            table_node_mapping,
            index_cls=VectorStoreIndex,
        )

    # 创建SQL查询工具
    sql_query_engine = SQLTableRetrieverQueryEngine(sql_database,
                                                    sql_obj_index.as_retriever(similarity_top_k=top_k),
                                                    verbose=True,
                                                    )
    return sql_query_engine

# Create a summary query engine
def create_summary_query_engine(index, top_k=3, use_reranker=False, filters=None):
    summary_index = SummaryIndex(index.vector_store.get_nodes(node_ids=None))
    summary_query_engine = summary_index.as_query_engine(
    response_mode=ResponseMode.TREE_SUMMARIZE,
    use_async=True,
    streaming=True,
    )
    return summary_query_engine

# Create a query engine
def create_query_engine(index, top_k=3, use_reranker=False, filters=None):
    # 创建向量检索查询工具
    postprocess = None
    if use_reranker:
        postprocess = get_node_postprocessors()

    query_engine = RetrieverQueryEngine.from_args(
        get_Retriever(index,
                      similarity_top_k=top_k,
                      filters=filters),
        text_qa_template=text_qa_template,
        refine_template=refine_template,
        summary_template = summary_template,
        simple_template = simple_template,
        node_postprocessors=postprocess,
        use_async=True,
        streaming=True,
    )

    return query_engine