nex_docus/backend/app/services/storage.py

"""
文件存储管理服务
"""
import os
import shutil
import uuid
import aiofiles
from pathlib import Path
from typing import List, Optional
from fastapi import HTTPException, UploadFile
from app.core.config import settings
from app.schemas.file import FileTreeNode


class StorageService:
    """文件存储服务类"""

    def __init__(self):
        # 获取 backend 目录的绝对路径（app/services 的父目录的父目录）
        backend_dir = Path(__file__).parent.parent.parent

        # 将配置中的路径转换为绝对路径
        self.projects_root = (backend_dir / settings.PROJECTS_PATH).resolve()
        self.temp_root = (backend_dir / settings.TEMP_PATH).resolve()

        # 确保根目录存在
        self.projects_root.mkdir(parents=True, exist_ok=True)
        self.temp_root.mkdir(parents=True, exist_ok=True)

    def get_secure_path(self, storage_key: str, relative_path: str = "") -> Path:
        """
        获取安全的文件路径（防止路径穿越攻击）

        Args:
            storage_key: 项目 UUID
            relative_path: 相对路径

        Returns:
            Path: 安全的绝对路径

        Raises:
            HTTPException: 非法路径访问
        """
        # 项目根目录
        project_root = self.projects_root / storage_key
        project_root = project_root.resolve()

        # 目标路径
        if relative_path:
            target_path = (project_root / relative_path).resolve()
        else:
            target_path = project_root

        # 安全检查：目标路径必须在项目根目录下
        try:
            target_path.relative_to(project_root)
        except ValueError:
            raise HTTPException(status_code=403, detail="非法路径访问")

        return target_path

    def create_project_structure(self, storage_key: str) -> None:
        """
        创建项目文件夹结构

        Args:
            storage_key: 项目 UUID
        """
        project_root = self.projects_root / storage_key
        project_root.mkdir(parents=True, exist_ok=True)

        # 创建 _assets 目录
        assets_dir = project_root / "_assets"
        assets_dir.mkdir(exist_ok=True)
        (assets_dir / "images").mkdir(exist_ok=True)
        (assets_dir / "files").mkdir(exist_ok=True)

        # 创建默认 README.md
        readme_path = project_root / "README.md"
        if not readme_path.exists():
            with open(readme_path, "w", encoding="utf-8") as f:
                f.write("# 项目首页\n\n欢迎使用 NEX Docus！\n")

    def generate_tree(self, path: Path, relative_root: str = "") -> List[FileTreeNode]:
        """
        生成目录树结构

        Args:
            path: 目录路径
            relative_root: 相对根路径

        Returns:
            List[FileTreeNode]: 目录树节点列表
        """
        if not path.exists() or not path.is_dir():
            return []

        tree = []
        try:
            # 获取所有文件和文件夹，按类型和名称排序
            items = sorted(
                path.iterdir(),
                key=lambda x: (not x.is_dir(), x.name.lower())
            )

            for item in items:
                # 跳过隐藏文件、特殊目录和 _assets 目录
                if item.name.startswith(".") or item.name == "_assets":
                    continue

                rel_path = str(Path(relative_root) / item.name) if relative_root else item.name

                node_data = {
                    "title": item.name,
                    "key": rel_path,
                    "isLeaf": item.is_file(),
                }

                if item.is_dir():
                    node_data["children"] = self.generate_tree(item, rel_path)

                tree.append(FileTreeNode(**node_data))

        except PermissionError:
            raise HTTPException(status_code=403, detail="权限不足")

        return tree

    async def read_file(self, file_path: Path) -> str:
        """
        读取文件内容

        Args:
            file_path: 文件路径

        Returns:
            str: 文件内容

        Raises:
            HTTPException: 文件不存在或无法读取
        """
        if not file_path.exists():
            raise HTTPException(status_code=404, detail="文件不存在")

        if not file_path.is_file():
            raise HTTPException(status_code=400, detail="不是有效的文件")

        try:
            async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
                content = await f.read()
            return content
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"文件读取失败: {str(e)}")

    async def write_file(self, file_path: Path, content: str) -> None:
        """
        写入文件内容

        Args:
            file_path: 文件路径
            content: 文件内容

        Raises:
            HTTPException: 写入失败
        """
        try:
            # 确保父目录存在
            file_path.parent.mkdir(parents=True, exist_ok=True)

            async with aiofiles.open(file_path, "w", encoding="utf-8") as f:
                await f.write(content)
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"文件写入失败: {str(e)}")

    async def delete_file(self, file_path: Path) -> None:
        """
        删除文件或文件夹

        Args:
            file_path: 文件路径

        Raises:
            HTTPException: 删除失败
        """
        if not file_path.exists():
            raise HTTPException(status_code=404, detail="文件不存在")

        try:
            if file_path.is_dir():
                shutil.rmtree(file_path)
            else:
                file_path.unlink()
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"删除失败: {str(e)}")

    async def rename_file(self, old_path: Path, new_path: Path) -> None:
        """
        重命名文件或文件夹

        Args:
            old_path: 旧路径
            new_path: 新路径

        Raises:
            HTTPException: 重命名失败
        """
        if not old_path.exists():
            raise HTTPException(status_code=404, detail="文件不存在")

        if new_path.exists():
            raise HTTPException(status_code=400, detail="目标路径已存在")

        try:
            old_path.rename(new_path)
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"重命名失败: {str(e)}")

    async def create_directory(self, dir_path: Path) -> None:
        """
        创建目录

        Args:
            dir_path: 目录路径

        Raises:
            HTTPException: 创建失败
        """
        if dir_path.exists():
            raise HTTPException(status_code=400, detail="目录已存在")

        try:
            dir_path.mkdir(parents=True, exist_ok=True)
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"创建目录失败: {str(e)}")

    async def upload_file(
        self,
        storage_key: str,
        file: UploadFile,
        subfolder: str = "images"
    ) -> dict:
        """
        上传文件到项目资源目录

        Args:
            storage_key: 项目 UUID
            file: 上传的文件
            subfolder: 子文件夹（images 或 files）

        Returns:
            dict: 文件信息

        Raises:
            HTTPException: 上传失败
        """
        # 生成唯一文件名
        file_ext = Path(file.filename).suffix
        unique_filename = f"{uuid.uuid4().hex}{file_ext}"

        # 目标路径
        target_dir = self.get_secure_path(storage_key, f"_assets/{subfolder}")
        target_dir.mkdir(parents=True, exist_ok=True)
        target_path = target_dir / unique_filename

        try:
            # 保存文件
            async with aiofiles.open(target_path, "wb") as f:
                content = await file.read()
                await f.write(content)

            # 返回文件信息
            relative_path = f"_assets/{subfolder}/{unique_filename}"
            return {
                "filename": unique_filename,
                "original_filename": file.filename,
                "path": relative_path,
                "size": len(content),
            }
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"文件上传失败: {str(e)}")

    async def upload_document(
        self,
        storage_key: str,
        file: UploadFile,
        target_dir: str = "",
        allowed_extensions: list = None
    ) -> dict:
        """
        上传文档文件到项目指定目录

        Args:
            storage_key: 项目 UUID
            file: 上传的文件
            target_dir: 目标目录（相对路径，如 "docs" 或 "docs/manuals"）
            allowed_extensions: 允许的文件扩展名列表，如 [".pdf", ".docx"]

        Returns:
            dict: 文件信息

        Raises:
            HTTPException: 上传失败或文件类型不允许
        """
        # 验证文件扩展名
        file_ext = Path(file.filename).suffix.lower()
        if allowed_extensions and file_ext not in allowed_extensions:
            raise HTTPException(
                status_code=400,
                detail=f"不支持的文件类型: {file_ext}。允许的类型: {', '.join(allowed_extensions)}"
            )

        # 使用原始文件名（不添加时间戳）
        unique_filename = file.filename

        # 目标路径
        if target_dir:
            target_path = self.get_secure_path(storage_key, target_dir)
        else:
            target_path = self.get_secure_path(storage_key)

        target_path.mkdir(parents=True, exist_ok=True)
        file_path = target_path / unique_filename

        try:
            # 保存文件
            async with aiofiles.open(file_path, "wb") as f:
                content = await file.read()
                await f.write(content)

            # 返回文件信息
            if target_dir:
                relative_path = f"{target_dir}/{unique_filename}"
            else:
                relative_path = unique_filename

            return {
                "filename": unique_filename,
                "original_filename": file.filename,
                "path": relative_path,
                "size": len(content),
            }
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"文件上传失败: {str(e)}")


# 创建全局实例
storage_service = StorageService()