main.py

#!/usr/bin/env python3
"""
AI 垃圾分类助手 - 环保小能手
主入口程序 (CLI)
"""

import argparse
from pathlib import Path
from config import DATASET_DIR, MODEL_DIR, MODEL_PATH


def cmd_train(args):
    from train import train
    train(args)


def cmd_predict(args):
    from predict import GarbageClassifier
    from knowledge import get_class_info

    classifier = GarbageClassifier(model_path=args.model_path)
    image_path = Path(args.image)
    if not image_path.exists():
        print(f"✗ 图片不存在: {image_path}")
        return

    print(f"正在识别: {image_path.name}")
    results = classifier.predict(str(image_path), top_k=args.top_k)

    print(f"\n{'='*40}")
    print("识别结果:")
    print(f"{'='*40}")
    for i, r in enumerate(results, 1):
        pct = r["confidence"] * 100
        bar = "█" * int(pct // 5) + "░" * (20 - int(pct // 5))
        print(f"  {i}. [{bar}] {r['class_name_cn']} ({pct:.1f}%)")

    if args.detail:
        info = get_class_info(results[0]["class_name"])
        if info:
            print(f"\n{'='*40}")
            print(f"垃圾分类指南 - {info['name_cn']}")
            print(f"分类: {info['category']}")
            print(f"{'='*40}")
            print(info["disposal"])
            print(f"\n💡 {info['fun_fact']}")


def cmd_query(args):
    from knowledge import search_knowledge

    results = search_knowledge(args.keyword)
    if not results:
        print(f"未找到与 '{args.keyword}' 相关的知识")
        return

    for _, info in results:
        print(f"\n{'='*50}")
        print(f"{info['name_cn']} | 类别: {info['category']}")
        print(f"{'='*50}")
        print(info["description"])
        print(f"\n📋 投放方法:")
        print(info["disposal"])
        print(f"\n💡 小贴士:")
        for tip in info["tips"]:
            print(f"  • {tip}")
        print(f"\n🎯 {info['fun_fact']}")
        print(f"⏱ 降解时间: {info['degradation_time']}")


def cmd_record(args):
    from database import Database
    from predict import GarbageClassifier
    from knowledge import get_class_info

    db = Database()
    user_id = db.register_user(args.username)

    classifier = GarbageClassifier(model_path=args.model_path)
    image_path = Path(args.image)
    if not image_path.exists():
        print(f"✗ 图片不存在: {image_path}")
        return

    results = classifier.predict(str(image_path))
    best = results[0]
    points = db.add_record(user_id, best["class_name"], best["confidence"])

    print(f"\n✓ 已记录! {best['class_name_cn']} (置信度: {best['confidence']*100:.1f}%)")
    print(f"  +{points} 环保积分!")

    info = get_class_info(best["class_name"])
    if info:
        print(f"\n📋 投放提示: {info['disposal'].split(chr(10))[0]}")


def cmd_stats(args):
    from database import Database
    from knowledge import KNOWLEDGE_BASE

    db = Database()
    user = db.get_user(args.username)
    if not user:
        print(f"用户 '{args.username}' 不存在，请先使用 record 命令")
        return

    stats = db.get_user_stats(user["id"])
    print(f"\n{'='*50}")
    print(f"  环保统计 - {stats['username']}")
    print(f"{'='*50}")
    print(f"  总分类次数: {stats['total_classifications']}")
    print(f"  总环保积分: {stats['total_points']}")
    print(f"  今日分类: {stats['today']['count']} 次")
    print(f"  今日积分: {stats['today']['points']}")

    if stats["class_distribution"]:
        print(f"\n  各类别分类统计:")
        for item in stats["class_distribution"]:
            cn = KNOWLEDGE_BASE.get(item["predicted_class"], {}).get("name_cn", item["predicted_class"])
            print(f"    • {cn}: {item['count']} 次")

    if stats["recent_records"]:
        print(f"\n  最近记录:")
        for r in stats["recent_records"][:5]:
            cn = KNOWLEDGE_BASE.get(r["predicted_class"], {}).get("name_cn", r["predicted_class"])
            print(f"    • {cn} | 积分: +{r['points']} | {r['created_at']}")


def cmd_leaderboard(args):
    from database import Database

    db = Database()
    leaders = db.get_leaderboard(args.limit)

    if not leaders:
        print("暂无环保数据，快去分类吧！")
        return

    print(f"\n{'='*50}")
    print("  🏆 环保积分排行榜")
    print(f"{'='*50}")
    print(f"  {'排名':>4} {'用户名':<15} {'积分':<8} {'分类次数':<8}")
    print(f"  {'-'*35}")
    badges = ["🥇", "🥈", "🥉"]
    for i, u in enumerate(leaders, 1):
        badge = badges[i - 1] if i <= 3 else "  "
        print(f"  {badge} {i:<2} {u['username']:<15} {u['total_points']:<8} {u['total_classifications']:<8}")


def cmd_web(args):
    from app.api import start_api
    start_api(host=args.host, port=args.port)


def cmd_webui(args):
    from webui import launch_gradio
    launch_gradio(server_port=args.port)


def main():
    parser = argparse.ArgumentParser(
        description="AI 垃圾分类助手 - 环保小能手",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  %(prog)s train --data-dir dataset/trashnet
  %(prog)s predict image.jpg
  %(prog)s query 塑料
  %(prog)s record --username 小明 image.jpg
  %(prog)s stats --username 小明
  %(prog)s leaderboard
  %(prog)s web          # 启动 API 服务 (供小程序调用)
  %(prog)s webui        # 启动 Gradio 网页界面
        """,
    )
    sub = parser.add_subparsers(dest="command", help="可用命令")

    p = sub.add_parser("train", help="训练分类模型")
    p.add_argument("--data-dir", default=str(DATASET_DIR))
    p.add_argument("--model-dir", default=str(MODEL_DIR))
    p.add_argument("--epochs", type=int, default=30)
    p.add_argument("--batch-size", type=int, default=32)
    p.add_argument("--lr", type=float, default=0.001)

    p = sub.add_parser("predict", help="分类垃圾图片")
    p.add_argument("image")
    p.add_argument("--model-path", default=str(MODEL_PATH))
    p.add_argument("--top-k", type=int, default=3)
    p.add_argument("--no-detail", dest="detail", action="store_false", default=True)

    p = sub.add_parser("query", help="查询垃圾分类知识")
    p.add_argument("keyword")

    p = sub.add_parser("record", help="分类并记录积分")
    p.add_argument("image")
    p.add_argument("--username", default="default")
    p.add_argument("--model-path", default=str(MODEL_PATH))

    p = sub.add_parser("stats", help="查看个人统计")
    p.add_argument("--username", default="default")

    p = sub.add_parser("leaderboard", help="查看排行榜")
    p.add_argument("--limit", type=int, default=10)

    p = sub.add_parser("web", help="启动 API 服务")
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--port", type=int, default=8000)

    p = sub.add_parser("webui", help="启动 Gradio 网页界面")
    p.add_argument("--port", type=int, default=7860)

    args = parser.parse_args()

    cmds = {
        "train": cmd_train,
        "predict": cmd_predict,
        "query": cmd_query,
        "record": cmd_record,
        "stats": cmd_stats,
        "leaderboard": cmd_leaderboard,
        "web": cmd_web,
        "webui": cmd_webui,
    }
    fn = cmds.get(args.command)
    if fn:
        fn(args)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()