from server.chat.check_language import check_language, get_supported_languages
from server.chat.chat_comparison import chat_comparison_test
from server.chat.gen_title import gen_title
from server.chat.relevant_articles import relevant_articles
from server.chat.self_kb_chat import self_kb_chat
from server.chat.stop import stop
import nltk
import sys
import os

from server.chat.chat_test import chat_test, get_image
from server.chat.gen_abstract import gen_abstract
from server.chat.gen_conclusion import gen_conclusion
from server.chat.gen_keywords import gen_keywords
from server.chat.gen_paragraph import gen_paragraph
from server.chat.knowledge_chat_test import knowledge_chat_test
from server.chat.translate import tarnslate_text
from server.chat.upload import upload_file
from server.chat.utils import download_self_doc
from server.chat.word_explain import word_explain
from server.chat.write_article import write_article
from server.knowledge_base.kb_doc_api import search_self_docs, upload_docs_new
from server.translator_service.main_api import cancel_task, download_result, get_progress, translate_file

sys.path.append(os.path.dirname(os.path.dirname(__file__)))

from configs import VERSION
from configs.model_config import NLTK_DATA_PATH
from configs.server_config import OPEN_CROSS_DOMAIN
import argparse
import uvicorn
from fastapi import Body
from fastapi.middleware.cors import CORSMiddleware
from starlette.responses import RedirectResponse
from server.chat.chat import chat
from server.chat.search_engine_chat import search_engine_chat
from server.chat.completion import completion
from server.custom.chapter_overview import chapter_overview
from server.custom.article_overview import article_overview
from server.custom.abstract_search import abstract_search
from server.custom.paper_translation import paper_translation
from server.chat.feedback import chat_feedback
from server.embeddings_api import embed_texts_endpoint
from server.llm_api import (list_running_models, list_config_models,
                            change_llm_model, stop_llm_model,
                            get_model_config, list_search_engines)
from server.utils import (BaseResponse, ListResponse, FastAPI, MakeFastAPIOffline, get_server_configs, get_prompt_template)
from typing import List, Literal
from server.chat.rewrite import(
    # con_rewrite,
    # exp_write,
    # abb_write,
    formal_style,
    party_style,
    col_style,
    chi_to_ens,
    ens_to_chi
)
from server.chat.con_rewrite import con_rewrite
from server.chat.exp_rewrite import exp_rewrite
from server.chat.abb_rewrite import abb_rewrite
from server.chat.rew_rewrite import rew_rewrite
from server.chat.sentence_reference import sentence_reference
from contextlib import asynccontextmanager
from server.translator_service.task_manager import TaskManager

nltk.data.path = [NLTK_DATA_PATH] + nltk.data.path



async def document():
    return RedirectResponse(url="/docs")

@asynccontextmanager
async def lifespan(app: FastAPI):
    tm = TaskManager(translate_fn=translate_file)
    tm.start()
    app.state.tm = tm
    # 2. 手动执行所有注册的 startup 钩子（包括启动器注入的）
    for fn in app.router.on_startup:
        await fn()
    yield
    tm.shutdown()
    for fn in app.router.on_shutdown:
        await fn()

def create_app(run_mode: str = None):
    app = FastAPI(
        title="Langchain-Chatchat API Server",
        version=VERSION,
        lifespan=lifespan,
    )
    MakeFastAPIOffline(app)
    # asyncio.create_task(lifespans)
    # Add CORS middleware to allow all origins
    # 在config.py中设置OPEN_DOMAIN=True，允许跨域
    # set OPEN_DOMAIN=True in config.py to allow cross-domain
    if OPEN_CROSS_DOMAIN:
        app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )
    mount_app_routes(app, run_mode=run_mode)
    return app


def mount_app_routes(app: FastAPI, run_mode: str = None):
    app.get("/",
            response_model=BaseResponse,
            summary="swagger 文档")(document)
    
    app.get("/chat/get_image",
            tags=["Chat"],
            summary="获取图片",
            )(get_image)

    app.get("/chat/get_self_file",
            tags=["Chat"],
            summary="获取个人知识库文件",
            )(download_self_doc)

    # Tag: Chat
    
    app.post("/chat/chat_comparison",
             tags=["Chat"],
             summary="生成文献对比报告"
             )(chat_comparison_test)
    
    app.post("/chat/chat",
             tags=["Chat"],
             summary="与llm模型对话(通过LLMChain)",
             )(chat_test)

    app.post("/chat",
             tags=["Chat"],
             summary="与llm模型对话(通过LLMChain)",
             )(chat)

    app.post("/chat/translate_text",
             tags=["Chat"],
             summary="翻译",
             )(tarnslate_text)

    app.post("/translate/translate_file",
             tags=["translate"],
             summary="文件翻译",
             )(translate_file)

    app.get("/translate/download_file",
             tags=["translate"],
             summary="下载译文",
             )(download_result)

    app.get("/translate/progress",
             tags=["translate"],
             summary="获取翻译任务进度",
             )(get_progress)
    
    app.post("/translate/cancel",
            tags=["translate"],
            summary="取消翻译任务",
            )(cancel_task)

    app.post("/chat/check_language",
             tags=["Chat"],
             summary="语种检测接口",
             )(check_language)
    
    app.get("/chat/get_language",
             tags=["Chat"],
             summary="获取当前支持的语种",
             )(get_supported_languages)
    
        # Tag: Chat
    app.post("/chat/stop",
             tags=["Chat"],
             summary="中断模型请求",
             )(stop)
    app.post("/chat/upload_Allfile",
             tags=["Chat"],
             summary="上传文件",
             )(upload_file)
    app.post("/chat/outlines",
             tags=["Chat"],
             summary="与llm模型对话生成大纲(通过LLMChain)",
             )(knowledge_chat_test)
    
    app.post("/chat/finsh_outlines",
             tags=["Chat"],
             summary="与llm模型对话生成全文(通过LLMChain)",
             )(write_article)
    
    app.post("/chat/chat_test",
             tags=["Chat"],
             summary="与llm模型对话(通过LLMChain)",
             )(chat_test)

    app.post("/chat/search_engine_chat",
             tags=["Chat"],
             summary="与搜索引擎对话",
             )(search_engine_chat)

    app.post("/chat/feedback",
             tags=["Chat"],
             summary="返回llm模型对话评分",
             )(chat_feedback)

    app.post("/chat/gen_title",
             tags=["Chat"],
             summary="生成当前对话的标题",
             )(gen_title)

    app.post("/rewrite/con_rewrite",
            tags=["Write"],
            summary="续写文本",
            )(con_rewrite)

    app.post("/rewrite/exp_write",
            tags=["Write"],
            summary="扩写文本",
            )(exp_rewrite)
    
    app.post("/rewrite/abb_write",
            tags=["Write"],
            summary="缩写文本",
            )(abb_rewrite)
    
    app.post("/rewrite/rew_rewrite",
            tags=["Write"],
            summary="重写文本",
            )(rew_rewrite)
    
    app.post("/sentence_reference",
            tags=["Write"],
            summary="好句子提示",
            )(sentence_reference)

    app.post("/gen_abstract",
            tags=["Write"],
            summary="摘要生成",
            )(gen_abstract)
    
    app.post("/gen_conclusion",
            tags=["Write"],
            summary="结论生成",
            )(gen_conclusion)
    
    app.post("/gen_keywords",
            tags=["Read"],
            summary="关键词生成",
            )(gen_keywords)

    app.post("/gen_paragraph",
            tags=["Read"],
            summary="章节速览",
            )(gen_paragraph)    
    
    app.post("/word_explain",
            tags=["Read"],
            summary="名词解释",
            )(word_explain)   
    
    app.post("/relevant_articles",
            tags=["Read"],
            summary="相关文献",
            )(relevant_articles)   
    #新功能接口
#     app.post("/rewrite/con_rewrite",
#             tags=["IastStrategy"],
#             summary="续写文本",
#             )(con_rewrite)
#
#     app.post("/rewrite/rewrite",
#             tags=["IastStrategy"],
#             summary="改写",
#             )(rewrite)
    # app.post("/rewrite/exp_write",
    #         tags=["IastStrategy"],
    #         summary="扩写",
    #         )(exp_write)

    # app.post("/rewrite/abb_write",
    #         tags=["IastStrategy"],
    #         summary="缩写",
    #         )(abb_write)

#     app.post("/rewrite/embellish",
#             tags=["IastStrategy"],
#             summary="润色",
#             )(embellish)

    app.post("/rewrite/formal_style",
            tags=["IastStrategy"],
            summary="正式风格",
            )(formal_style)

    app.post("/rewrite/party_style",
            tags=["IastStrategy"],
            summary="党政风格",
            )(party_style)

    app.post("/rewrite/col_style",
            tags=["IastStrategy"],
            summary="口语风格",
            )(col_style)

    app.post("/rewrite/chi_to_ens",
            tags=["IastStrategy"],
            summary="中译英",
            )(chi_to_ens)

    app.post("/rewrite/ens_to_chi",
            tags=["IastStrategy"],
            summary="英译中",
            )(ens_to_chi)
    
    # 知识库相关接口
    mount_knowledge_routes(app)
    # 摘要相关接口
    mount_filename_summary_routes(app)

    # LLM模型相关接口
    app.post("/llm_model/list_running_models",
             tags=["LLM Model Management"],
             summary="列出当前已加载的模型",
             )(list_running_models)

    app.post("/llm_model/list_config_models",
             tags=["LLM Model Management"],
             summary="列出configs已配置的模型",
             )(list_config_models)

    app.post("/llm_model/get_model_config",
             tags=["LLM Model Management"],
             summary="获取模型配置（合并后）",
             )(get_model_config)

    app.post("/llm_model/stop",
             tags=["LLM Model Management"],
             summary="停止指定的LLM模型（Model Worker)",
             )(stop_llm_model)

    app.post("/llm_model/change",
             tags=["LLM Model Management"],
             summary="切换指定的LLM模型（Model Worker)",
             )(change_llm_model)

    # 服务器相关接口
    app.post("/server/configs",
             tags=["Server State"],
             summary="获取服务器原始配置信息",
             )(get_server_configs)

    app.post("/server/list_search_engines",
             tags=["Server State"],
             summary="获取服务器支持的搜索引擎",
             )(list_search_engines)
        
    @app.post("/server/get_prompt_template",
              tags=["Server State"],
              summary="获取服务区配置的 prompt 模板")

    def get_server_prompt_template(
            type: Literal[
            	"llm_chat",
            	"knowledge_base_chat",
                "report_chat",
            	"search_engine_chat",
            	"agent_chat"
            ] = Body("llm_chat",
                     description="模板类型，可选值：llm_chat，knowledge_base_chat，search_engine_chat，agent_chat"),
            name: str = Body("default", description="模板名称"),
    ) -> str:
        return get_prompt_template(type=type, name=name)

    # 其它接口
    app.post("/other/completion",
             tags=["Other"],
             summary="要求llm模型补全(通过LLMChain)",
             )(completion)

    app.post("/other/embed_texts",
             tags=["Other"],
             summary="将文本向量化，支持本地模型和在线模型",
             )(embed_texts_endpoint)
    app.post("/knowledge_base/chapter_overview",
             tags=["Other"],
             summary="文件速览"
             )(chapter_overview)
    app.post("/knowledge_base/abstract_search",
             tags=["Other"],
             summary="相似摘要搜索"
             )(abstract_search)
    app.post("/knowledge_base/article_overview",
             tags=["Other"],
             summary="文件综述"
             )(article_overview)
    app.post("/knowledge_base/paper_translation",
             tags=["Other"],
             summary="论文翻译"
             )(paper_translation)

def mount_knowledge_routes(app: FastAPI):
    from server.chat.knowledge_base_chat import knowledge_base_chat
    from server.chat.knowledge_base_chat_old import knowledge_base_chat_old
    from server.chat.report_chat import report_chat
    from server.chat.file_chat import upload_temp_docs, file_chat
    from server.chat.agent_chat import agent_chat
    from server.knowledge_base.kb_api import list_kbs, create_kb, delete_kb
    from server.knowledge_base.kb_doc_api import (list_files, upload_docs, delete_docs,
                                                  update_docs, download_doc, recreate_vector_store,
                                                  search_docs, DocumentWithVSId, update_info,
                                                  update_docs_by_id, )

    app.post("/chat/knowledge_base_chat",
             tags=["Chat"],
             summary="与知识库对话")(knowledge_base_chat)

    app.post("/chat/self_kb_chat",
             tags=["Chat"],
             summary="与个人知识库对话")(self_kb_chat)

    app.post("/chat/knowledge_base_chat_old",
             tags=["Chat"],
             summary="旧版与知识库对话")(knowledge_base_chat_old)

    app.post("/chat/report_chat",
             tags=["Chat"],
             summary="与报告知识库对话")(report_chat)
    
    app.post("/chat/file_chat",
             tags=["Knowledge Base Management"],
             summary="文件对话"
             )(file_chat)

    app.post("/chat/agent_chat",
             tags=["Chat"],
             summary="与agent对话")(agent_chat)

    # Tag: Knowledge Base Management
    app.get("/knowledge_base/list_knowledge_bases",
            tags=["Knowledge Base Management"],
            response_model=ListResponse,
            summary="获取知识库列表")(list_kbs)

    app.post("/knowledge_base/create_knowledge_base",
             tags=["Knowledge Base Management"],
             response_model=BaseResponse,
             summary="创建知识库"
             )(create_kb)

    app.post("/knowledge_base/delete_knowledge_base",
             tags=["Knowledge Base Management"],
             response_model=BaseResponse,
             summary="删除知识库"
             )(delete_kb)

    app.get("/knowledge_base/list_files",
            tags=["Knowledge Base Management"],
            response_model=ListResponse,
            summary="获取知识库内的文件列表"
            )(list_files)

    app.post("/knowledge_base/search_docs",
             tags=["Knowledge Base Management"],
             response_model=List[DocumentWithVSId],
             summary="搜索知识库"
             )(search_docs)
    
    app.post("/knowledge_base/search_self_docs",
             tags=["Knowledge Base Management"],
             response_model=List[DocumentWithVSId],
             summary="搜索个人知识库"
             )(search_self_docs)
    app.post("/knowledge_base/update_docs_by_id",
             tags=["Knowledge Base Management"],
             response_model=BaseResponse,
             summary="直接更新知识库文档"
             )(update_docs_by_id)

    app.post("/knowledge_base/upload_docs",
             tags=["Knowledge Base Management"],
             response_model=BaseResponse,
             summary="上传文件到知识库，并/或进行向量化"
             )(upload_docs)
    

    app.post("/knowledge_base/upload_docs_new",
             tags=["Knowledge Base Management"],
             response_model=BaseResponse,
             summary="上传文件到知识库，并/或进行向量化，并获取解析结果"
             )(upload_docs_new)

    app.post("/knowledge_base/delete_docs",
             tags=["Knowledge Base Management"],
             response_model=BaseResponse,
             summary="删除知识库内指定文件"
             )(delete_docs)

    app.post("/knowledge_base/update_info",
             tags=["Knowledge Base Management"],
             response_model=BaseResponse,
             summary="更新知识库介绍"
             )(update_info)
    app.post("/knowledge_base/update_docs",
             tags=["Knowledge Base Management"],
             response_model=BaseResponse,
             summary="更新现有文件到知识库"
             )(update_docs)

    app.get("/knowledge_base/download_doc",
            tags=["Knowledge Base Management"],
            summary="下载对应的知识文件")(download_doc)

    app.post("/knowledge_base/recreate_vector_store",
             tags=["Knowledge Base Management"],
             summary="根据content中文档重建向量库，流式输出处理进度。"
             )(recreate_vector_store)

    app.post("/knowledge_base/upload_temp_docs",
             tags=["Knowledge Base Management"],
             summary="上传文件到临时目录，用于文件对话。"
             )(upload_temp_docs)


def mount_filename_summary_routes(app: FastAPI):
    from server.knowledge_base.kb_summary_api import (summary_file_to_vector_store, recreate_summary_vector_store,
                                                      summary_doc_ids_to_vector_store)

    app.post("/knowledge_base/kb_summary_api/summary_file_to_vector_store",
             tags=["Knowledge kb_summary_api Management"],
             summary="单个知识库根据文件名称摘要"
             )(summary_file_to_vector_store)
    app.post("/knowledge_base/kb_summary_api/summary_doc_ids_to_vector_store",
             tags=["Knowledge kb_summary_api Management"],
             summary="单个知识库根据doc_ids摘要",
             response_model=BaseResponse,
             )(summary_doc_ids_to_vector_store)
    app.post("/knowledge_base/kb_summary_api/recreate_summary_vector_store",
             tags=["Knowledge kb_summary_api Management"],
             summary="重建单个知识库文件摘要"
             )(recreate_summary_vector_store)


def run_api(host, port, **kwargs):
    if kwargs.get("ssl_keyfile") and kwargs.get("ssl_certfile"):
        uvicorn.run(app,
                    host=host,
                    port=port,
                    ssl_keyfile=kwargs.get("ssl_keyfile"),
                    ssl_certfile=kwargs.get("ssl_certfile"),
                    )
    else:
        uvicorn.run(app, host=host, port=port)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(prog='langchain-ChatGLM',
                                     description='About langchain-ChatGLM, local knowledge based ChatGLM with langchain'
                                                 ' ｜ 基于本地知识库的 ChatGLM 问答')
    parser.add_argument("--host", type=str, default="0.0.0.0")
    parser.add_argument("--port", type=int, default=7861)
    parser.add_argument("--ssl_keyfile", type=str)
    parser.add_argument("--ssl_certfile", type=str)
    # 初始化消息
    args = parser.parse_args()
    args_dict = vars(args)

    app = create_app()

    run_api(host=args.host,
            port=args.port,
            ssl_keyfile=args.ssl_keyfile,
            ssl_certfile=args.ssl_certfile,
            )