main.py

"""
智能垃圾分类系统 - FastAPI 后端
Smart Garbage Classification System - FastAPI Backend
"""

import os
import sys
import tempfile
import logging
from pathlib import Path
from contextlib import asynccontextmanager

from fastapi import FastAPI, File, UploadFile, HTTPException, Body
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware

from model import GarbageClassifier
from knowledge_base import query_dna, get_knowledge_base

# 日志配置
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)

# 全局分类器（延迟加载）
classifier: GarbageClassifier = None

# 获取前端目录
FRONTEND_DIR = Path(__file__).parent / "frontend"

# 允许的图片格式
ALLOWED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB


@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    logger.info("🚀 智能垃圾分类系统启动中...")
    # 预加载 DNA 知识库
    kb = get_knowledge_base()
    logger.info(f"🧬 垃圾 DNA 知识库已加载 ({kb.count} 条记录)")
    yield
    logger.info("👋 系统已关闭")


app = FastAPI(
    title="智能垃圾分类系统",
    description="上传垃圾照片，AI自动识别垃圾分类",
    version="1.0.0",
    lifespan=lifespan,
)

# CORS - 允许手机等设备跨域访问
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


def get_classifier() -> GarbageClassifier:
    """获取分类器（延迟加载单例）"""
    global classifier
    if classifier is None:
        logger.info("正在初始化垃圾分类模型...")
        classifier = GarbageClassifier()
    return classifier


# 服务前端静态文件
if FRONTEND_DIR.exists():
    app.mount("/static", StaticFiles(directory=str(FRONTEND_DIR)), name="static")


@app.get("/", response_class=HTMLResponse)
async def index():
    """返回主页面"""
    index_path = FRONTEND_DIR / "index.html"
    if index_path.exists():
        with open(index_path, "r", encoding="utf-8") as f:
            return f.read()
    return HTMLResponse("<h1>智能垃圾分类系统</h1><p>前端文件未找到</p>", status_code=200)


@app.get("/api/health")
async def health():
    """健康检查"""
    return {
        "status": "ok",
        "model_loaded": classifier is not None,
        "frontend_exists": FRONTEND_DIR.exists(),
    }


@app.post("/api/classify")
async def classify_image(file: UploadFile = File(...)):
    """
    上传图片并进行垃圾分类识别

    - **file**: 图片文件（jpg, png, webp, bmp），最大 10MB
    """
    # 验证文件类型
    if not file.content_type or not file.content_type.startswith("image/"):
        raise HTTPException(
            status_code=400,
            detail="请上传图片文件（支持 JPG、PNG、WebP、BMP 格式）"
        )

    # 验证文件扩展名
    ext = os.path.splitext(file.filename or "image.jpg")[1].lower()
    if ext not in ALLOWED_EXTENSIONS and ext:
        raise HTTPException(
            status_code=400,
            detail=f"不支持的图片格式: {ext}。支持: {', '.join(ALLOWED_EXTENSIONS)}"
        )

    # 读取文件内容
    content = await file.read()
    if len(content) > MAX_FILE_SIZE:
        raise HTTPException(
            status_code=400,
            detail=f"图片太大（{len(content)/1024/1024:.1f}MB），请上传小于 10MB 的图片"
        )

    if len(content) == 0:
        raise HTTPException(status_code=400, detail="文件为空")

    # 保存到临时文件
    suffix = ext if ext else ".jpg"
    tmp_path = None
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            tmp.write(content)
            tmp_path = tmp.name

        # 分类
        clf = get_classifier()
        result = clf.classify(tmp_path)
        return JSONResponse(result)

    except RuntimeError as e:
        if "transformers" in str(e) or "model" in str(e).lower():
            raise HTTPException(
                status_code=503,
                detail=f"模型加载失败: {str(e)}。请确保已安装依赖: pip install -r requirements.txt"
            )
        raise HTTPException(status_code=500, detail=f"分类失败: {str(e)}")
    except Exception as e:
        logger.error(f"分类过程中出错: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"分类失败: {str(e)}")
    finally:
        # 清理临时文件
        if tmp_path and os.path.exists(tmp_path):
            try:
                os.unlink(tmp_path)
            except Exception as e:
                logger.warning(f"清理临时文件失败: {e}")


@app.post("/api/classify-text")
async def classify_text(data: dict = Body(...)):
    """
    文本垃圾分类 - 根据物品名称判断垃圾分类
    """
    item_name = data.get("item", "").strip()
    if not item_name:
        raise HTTPException(status_code=400, detail="请提供物品名称（item 参数）")

    clf = get_classifier()
    result = clf.classify_text(item_name)
    return JSONResponse(result)


@app.get("/api/item-dna")
async def get_item_dna(item: str = None, item_zh: str = None):
    """查询物品的 DNA 知识

    - **item**: 英文物品名（模型预测标签）
    - **item_zh**: 中文物品名
    """
    if not item and not item_zh:
        return {
            "dna_entries": get_knowledge_base().count,
            "detail": "请提供 item 或 item_zh 参数进行查询"
        }

    dna = query_dna(item, item_zh)
    if dna:
        return {"found": True, "dna": dna}
    return {"found": False, "detail": f"未找到「{item_zh or item}」的 DNA 信息"}


@app.get("/api/daily-challenge")
async def get_daily_challenge():
    """获取今日每日挑战题目（10 件物品）"""
    from datetime import date
    clf = get_classifier()
    items = clf.generate_daily_challenge(date.today())
    return {
        "date": date.today().isoformat(),
        "items": items,
        "total": len(items),
    }


@app.post("/api/daily-challenge/submit")
async def submit_answer(data: dict = Body(...)):
    """提交单题答案，返回对错及趣味知识

    Body: { item_en: str, guess: str }
     - item_en: 物品英文标签
     - guess: 用户猜测的分类
    Returns: { correct: bool, correct_category: str, dna: {...} | null }
    """
    item_en = data.get("item_en", "")
    guess = data.get("guess", "")

    if not item_en or not guess:
        return {"error": "缺少 item_en 或 guess 参数"}

    clf = get_classifier()
    correct_category = clf.label_to_category.get(item_en, "其他垃圾")
    is_correct = guess == correct_category

    # 查询 DNA
    item_zh = clf._get_label_zh(item_en)
    dna = query_dna(item_en, item_zh)

    return {
        "correct": is_correct,
        "correct_category": correct_category,
        "item_zh": item_zh,
        "dna": dna,
    }


@app.get("/api/categories")
async def get_categories():
    """获取分类信息"""
    clf = get_classifier()
    return {
        "categories": [
            {
                "name": name,
                **info
            }
            for name, info in clf.category_info.items()
        ]
    }


# 错误处理
@app.exception_handler(404)
async def not_found(request, exc):
    return JSONResponse(
        status_code=404,
        content={"detail": "请求的资源不存在"}
    )


@app.exception_handler(500)
async def server_error(request, exc):
    return JSONResponse(
        status_code=500,
        content={"detail": "服务器内部错误，请稍后重试"}
    )


if __name__ == "__main__":
    import uvicorn

    host = os.getenv("HOST", "0.0.0.0")
    port = int(os.getenv("PORT", "8000"))

    print("=" * 50)
    print("  🌱 智能垃圾分类系统")
    print("=" * 50)
    print(f"  本地访问: http://localhost:{port}")
    print(f"  局域网访问: http://<本机IP>:{port}")
    print(f"  上传图片即可识别垃圾类别")
    print("=" * 50)

    uvicorn.run(app, host=host, port=port, log_level="info")