# utils/file_handler.py
import aiofiles
import hashlib
import re
import uuid
from pathlib import Path
from fastapi import UploadFile
from typing import Tuple, Dict, Any


class FileHandler:
    def __init__(self, upload_dir: str = "uploads", max_file_size: int = 50 * 1024 * 1024):
        self.upload_dir = Path(upload_dir)
        self.upload_dir.mkdir(exist_ok=True)
        self.max_file_size = max_file_size

    def _sanitize_filename(self, filename: str) -> str:
        original = Path(filename or "upload.step").name
        suffix = Path(original).suffix.lower()
        stem = Path(original).stem or "upload"
        safe_stem = re.sub(r"[^A-Za-z0-9._-]+", "_", stem).strip("._-") or "upload"
        if suffix not in {".stp", ".step"}:
            suffix = ".step"
        return f"{safe_stem}{suffix}"

    @staticmethod
    def _looks_like_step(content: bytes) -> bool:
        if not content:
            return False
        head = content[:4096].decode("utf-8", errors="ignore").upper()
        return (
            "ISO-10303-21" in head
            or "HEADER;" in head
            or "FILE_SCHEMA" in head
            or "DATA;" in head
        )

    async def save_uploaded_file(self, file: UploadFile) -> Tuple[Path, int, Dict[str, Any]]:
        """保存上传的文件并返回安全元数据"""
        content = await file.read()

        if not content:
            raise ValueError("上传文件为空")
        if len(content) > self.max_file_size:
            raise ValueError(f"上传文件过大，限制 {self.max_file_size // (1024 * 1024)}MB")
        if not self._looks_like_step(content):
            raise ValueError("文件内容不是有效的 STP/STEP 数据")

        safe_name = self._sanitize_filename(file.filename)
        unique_name = f"{uuid.uuid4().hex}_{safe_name}"
        file_path = self.upload_dir / unique_name

        async with aiofiles.open(file_path, "wb") as f:
            await f.write(content)

        metadata = {
            "original_filename": file.filename,
            "safe_original_name": safe_name,
            "stored_filename": unique_name,
            "sha256": hashlib.sha256(content).hexdigest(),
        }

        return file_path, len(content), metadata

    def cleanup_file(self, file_path: Path):
        """清理文件"""
        try:
            if file_path.exists():
                file_path.unlink()
        except Exception as e:
            print(f"文件清理失败: {e}")