ProxyPool/app/api/routes/plugins.py

"""插件相关路由"""
import asyncio
from fastapi import APIRouter, Depends
from app.services.plugin_service import PluginService
from app.services.scheduler_service import SchedulerService
from app.services.task_service import task_service
from app.api.deps import get_plugin_service, get_scheduler_service
from app.api.common import success_response, error_response, format_plugin
from app.core.log import logger

router = APIRouter(prefix="/api/plugins", tags=["plugins"])


@router.get("")
async def list_plugins(service: PluginService = Depends(get_plugin_service)):
    plugins = await service.list_plugins()
    return success_response("获取插件列表成功", {"plugins": [format_plugin(p) for p in plugins]})


@router.put("/{plugin_id}/toggle")
async def toggle_plugin(
    plugin_id: str,
    request: dict,
    service: PluginService = Depends(get_plugin_service),
):
    enabled = request.get("enabled")
    if enabled is None:
        return error_response("缺少 enabled 参数", 400)

    success = await service.toggle_plugin(plugin_id, enabled)
    if not success:
        return error_response("插件不存在", 404)
    return success_response(
        f"插件 {plugin_id} 已{'启用' if enabled else '禁用'}",
        {"plugin_id": plugin_id, "enabled": enabled},
    )


@router.get("/{plugin_id}/config")
async def get_plugin_config(
    plugin_id: str,
    service: PluginService = Depends(get_plugin_service),
):
    config = await service.get_plugin_config(plugin_id)
    if config is None:
        return error_response("插件不存在", 404)
    return success_response("获取插件配置成功", {"plugin_id": plugin_id, "config": config})


@router.post("/{plugin_id}/config")
async def update_plugin_config(
    plugin_id: str,
    request: dict,
    service: PluginService = Depends(get_plugin_service),
):
    config = request.get("config", {})
    if not isinstance(config, dict):
        return error_response("config 必须是对象", 400)

    success = await service.update_plugin_config(plugin_id, config)
    if not success:
        return error_response("插件不存在或配置无效", 404)
    return success_response("保存插件配置成功", {"plugin_id": plugin_id, "config": config})


@router.post("/{plugin_id}/crawl")
async def crawl_plugin(
    plugin_id: str,
    plugin_service: PluginService = Depends(get_plugin_service),
    scheduler_service: SchedulerService = Depends(get_scheduler_service),
):
    plugin = plugin_service.get_plugin(plugin_id)
    if not plugin:
        return error_response("插件不存在", 404)

    task_id = task_service.create_task("crawl_plugin")

    async def _background_crawl():
        try:
            task_service.update_task(task_id, status="running")
            results = await plugin_service.run_plugin(plugin_id)
            if not results:
                task_service.update_task(
                    task_id,
                    status="completed",
                    message=f"插件 {plugin_id} 爬取完成，未获取到代理",
                    data={"plugin_id": plugin_id, "proxy_count": 0, "valid_count": 0},
                )
                return

            logger.info(f"Plugin {plugin_id} crawled {len(results)} proxies")
            scheduler_service.validation_queue.reset_stats()
            await scheduler_service.validation_queue.submit(results)

            try:
                await asyncio.wait_for(scheduler_service.validation_queue.drain(), timeout=5.0)
            except asyncio.TimeoutError:
                pass

            task_service.update_task(
                task_id,
                status="completed",
                message=f"插件 {plugin_id} 爬取并验证完成",
                data={
                    "plugin_id": plugin_id,
                    "proxy_count": len(results),
                    "valid_count": scheduler_service.validation_queue.valid_count,
                    "invalid_count": scheduler_service.validation_queue.invalid_count,
                },
            )
        except Exception as e:
            logger.error(f"Crawl plugin {plugin_id} failed: {e}")
            task_service.update_task(
                task_id, status="failed", message=f"插件爬取失败: {str(e)}"
            )

    asyncio.create_task(_background_crawl())
    return success_response(
        "爬取任务已启动", {"task_id": task_id, "plugin_id": plugin_id}
    )


@router.post("/crawl-all")
async def crawl_all(
    plugin_service: PluginService = Depends(get_plugin_service),
    scheduler_service: SchedulerService = Depends(get_scheduler_service),
):
    task_id = task_service.create_task("crawl_all")

    async def _background_crawl_all():
        try:
            task_service.update_task(task_id, status="running")
            results = await plugin_service.run_all_plugins()
            if not results:
                task_service.update_task(
                    task_id,
                    status="completed",
                    message="所有插件爬取完成，未获取到代理",
                    data={"total_crawled": 0, "valid_count": 0, "invalid_count": 0},
                )
                return

            logger.info(f"All plugins crawled {len(results)} unique proxies")
            scheduler_service.validation_queue.reset_stats()
            await scheduler_service.validation_queue.submit(results)

            try:
                await asyncio.wait_for(scheduler_service.validation_queue.drain(), timeout=5.0)
            except asyncio.TimeoutError:
                pass

            task_service.update_task(
                task_id,
                status="completed",
                message="所有插件爬取并验证完成",
                data={
                    "total_crawled": len(results),
                    "valid_count": scheduler_service.validation_queue.valid_count,
                    "invalid_count": scheduler_service.validation_queue.invalid_count,
                },
            )
        except Exception as e:
            logger.error(f"Crawl all failed: {e}")
            task_service.update_task(
                task_id, status="failed", message=f"批量爬取失败: {str(e)}"
            )

    asyncio.create_task(_background_crawl_all())
    return success_response("爬取任务已启动", {"task_id": task_id})