from fastapi import APIRouter, UploadFile, File, Form, HTTPException
import os
import shutil
import logging
from scripts.rag_build_query import build_user_index

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

router = APIRouter()
ALLOWED_SUFFIXES = {".txt", ".md", ".pdf", ".docx"}

@router.post("/upload")
def upload_user_file(user_id: str = Form(...), file: UploadFile = File(...)):
    filename = os.path.basename(file.filename)
    suffix = os.path.splitext(filename)[-1].lower()

    if suffix not in ALLOWED_SUFFIXES:
        raise HTTPException(status_code=400, detail="不支持的文件类型")

    # 创建用户文档目录
    user_doc_dir = os.path.join("docs", user_id)
    if not os.path.exists(user_doc_dir):  # 检查目录是否存在
        os.makedirs(user_doc_dir, exist_ok=True)  # 如果不存在，创建目录

    # 创建索引数据目录，确保路径存在
    index_data_dir = os.path.join("index_data", user_id)
    if not os.path.exists(index_data_dir):  # 检查目录是否存在
        os.makedirs(index_data_dir, exist_ok=True)  # 如果不存在，创建目录

    # 保存文件
    file_path = os.path.join(user_doc_dir, filename)
    try:
        with open(file_path, "wb") as f:
            shutil.copyfileobj(file.file, f)
        logger.info(f"文件已保存至 {file_path}")

        # 创建索引并保存 docstore.json
        docstore_path = os.path.join(index_data_dir, "docstore.json")
        if not os.path.exists(docstore_path):  # 检查文件是否存在
            with open(docstore_path, "w") as docstore_file:
                docstore_file.write("{}")  # 可根据需要保存实际内容

        # 重新构建用户索引
        build_user_index(user_id)
        logger.info(f"用户 {user_id} 的索引已重建")

    except Exception as e:
        logger.error(f"[UPLOAD ERROR] {e}")
        raise HTTPException(status_code=500, detail="索引构建失败")

    return {"status": "ok", "filename": filename}