from dotenv import load_dotenv
from llama_index.core.node_parser import SentenceSplitter

load_dotenv()

import logging
import os
import uvicorn
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import RedirectResponse
from app.api.routers.chat import chat_router
from app.api.routers.upload import file_upload_router
from app.settings import init_settings
from app.observability import init_observability
from fastapi.staticfiles import StaticFiles
from phoenix.trace import using_project

from llama_index.core.prompts.default_prompts import (
    DEFAULT_REFINE_PROMPT,
    DEFAULT_REFINE_TABLE_CONTEXT_PROMPT,
    DEFAULT_TEXT_QA_PROMPT,
    DEFAULT_TREE_SUMMARIZE_PROMPT,
)

logger = logging.getLogger("uvicorn")

DEFAULT_TEXT_QA_PROMPT.template = (
    "# 角色\n"
    "你是一名博微造价工程数据查询助手，专精于电力工程文件中的信息。"
    "你的职责是提供有关电力造价、造价编制软件、文件结构及相关数据的精准、客观的回答，"
    "如同直接从文件中提取的内容。\n"
    
    "## 技能\n"
    "### 技能 1: 数据查询与提供\n"
    "- 准确回答所有关于电力工程造价的相关问题。\n"
    "- 提供具体数据，如成本估算、材料清单、劳动力需求等。\n"
    "- 确保提供的信息严格基于工程文档中的记录。\n"
    
    "### 技能 2: 技术性解释\n"
    "- 解释造价工程中的技术术语和概念。\n"
    "- 为复杂的工程细节提供清晰易懂的说明。\n"
    
    "## 约束\n"
    "- 仅回答与电力工程造价文件相关的具体问题。\n"
    "- 不进行任何超出文件内容的猜测或假设。\n"
    "- 所有回答均基于文件内容，采用客观和技术性的语言。\n"
    "- 请基于这些信息回答问题。如果无法找到相关信息，请不要额外发散回答，不要回答多余的信息，只需要回答“我不知道这个问题的答案”。\n"
    "上下文信息如下。\n"
    "---------------------\n"
    "{context_str}\n"
    "---------------------\n"
    "请仅依据上下文信息回答问题。\n"
    "Query: {query_str}\n"
    "Answer: "
)


DEFAULT_REFINE_PROMPT.template = (
    "原始查询如下: {query_str}\n"
    "已提供的答案: {existing_answer}\n"
    "现在有更多的上下文信息，可能有助于改进答案。\n"
    "------------\n"
    "{context_msg}\n"
    "------------\n"
    "请根据新的上下文改进原始答案，以更好地回答查询。"
    "如果新的上下文信息无用，请保留原答案。\n"
    "改进后的答案: "
)


DEFAULT_TREE_SUMMARIZE_PROMPT.template = (
    "# 角色\n"
    "你是一名博微造价工程数据查询助手，专精于电力工程文件中的信息。"
    "你的职责是提供有关电力造价、造价编制软件、文件结构及相关数据的精准、客观的回答，"
    "如同直接从文件中提取的内容。\n"
    
    "## 技能\n"
    "### 技能 1: 数据查询与提供\n"
    "- 准确回答所有关于电力工程造价的相关问题。\n"
    "- 提供具体数据，如成本估算、材料清单、劳动力需求等。\n"
    "- 确保提供的信息严格基于工程文档中的记录。\n"
    
    "### 技能 2: 技术性解释\n"
    "- 解释造价工程中的技术术语和概念。\n"
    "- 为复杂的工程细节提供清晰易懂的说明。\n"
    
    "## 约束\n"
    "- 仅回答与电力工程造价文件相关的具体问题。\n"
    "- 不进行任何超出文件内容的猜测或假设。\n"
    "- 所有回答均基于文件内容，采用客观和技术性的语言。\n"
    "- 请基于这些信息回答问题。如果无法找到相关信息，请不要额外发散回答，不要回答多余的信息，只需要回答“我不知道这个问题的答案”。\n"
    "来自多个来源的上下文信息如下。\n"
    "---------------------\n"
    "{context_str}\n"
    "---------------------\n"
    "鉴于来自多个来源的信息而非先验知识， "
    "回答查询。\n"
    "Query: {query_str}\n"
    "Answer: "
)

usPrj = using_project(os.getenv("PHOENIX_PROJECT_NAME"))
usPrj.__enter__()

init_settings()
init_observability()

app = FastAPI()

environment = os.getenv("ENVIRONMENT", "dev")  # Default to 'development' if not set
if environment == "dev":
    logger.warning("Running in development mode - allowing CORS for all origins")
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

def mount_static_files(directory, path):
    if os.path.exists(directory):
        for dir, _, _ in os.walk(directory):
            relative_path = os.path.relpath(dir, directory)
            mount_path = path if relative_path == "." else f"{path}/{relative_path}"
            logger.info(f"Mounting static files '{dir}' at {mount_path}")
            app.mount(mount_path, StaticFiles(directory=dir), name=f"{dir}-static")

# Mount the data files to serve the file viewer
mount_static_files("data", "/api/files/data")
# Mount the output files from tools
mount_static_files("data_output", "/api/files/output")
app.include_router(chat_router, prefix="/api/chat")
app.include_router(file_upload_router, prefix="/api/chat/upload")

@app.get("/")
async def redirect_to_docs():
    return RedirectResponse(url="/docs")

SentenceSplitter
if __name__ == "__main__":
    app_host = os.getenv("APP_HOST", "0.0.0.0")
    app_port = int(os.getenv("APP_PORT", "8000"))
    reload = True if environment == "dev" else False
    reload = False
    uvicorn.run(app="main:app", host=app_host, port=app_port, reload=reload)

    #usPrj.__exit__()