import argparse
import asyncio
import logging
from typing import Any

from app.core.config import get_settings
from app.core.errors import RecoverableAgentError
from app.infrastructure.db import engine
from app.infrastructure.milvus_profile_projection import MilvusProfileProjection
from app.infrastructure.neo4j_profile_projection import Neo4jProfileProjection
from app.service.agent.bootstrap import get_memory_embedding_service
from app.service.model_gateway import DatabaseModelEndpointResolver
from app.worker.memory_sync_outbox_worker import MemorySyncOutboxWorker
from app.worker.offsite_mail_worker import OffsiteMailWorker
from app.worker.runtime import WorkerRuntime

logger = logging.getLogger(__name__)


async def serve(*, once: bool = False) -> None:
    settings = get_settings()
    runtime = WorkerRuntime(settings=settings)
    offsite_worker = OffsiteMailWorker(settings)
    neo4j_driver: Any | None = None
    milvus_client: Any | None = None
    memory_sync_worker: MemorySyncOutboxWorker | None = None
    memory_sync_handlers: dict[str, Any] = {}
    if settings.neo4j_password:
        from neo4j import AsyncGraphDatabase

        neo4j_driver = AsyncGraphDatabase.driver(
            settings.neo4j_uri,
            auth=(settings.neo4j_username, settings.neo4j_password),
        )
        memory_sync_worker = MemorySyncOutboxWorker({
            "neo4j": Neo4jProfileProjection(neo4j_driver).upsert,
        })
    else:
        logger.warning("Neo4j password not configured; profile projection remains pending")
    if settings.resolved_milvus_uri and settings.knowledge_embedding_endpoint_code:
        from pymilvus import AsyncMilvusClient  # type: ignore[import-untyped]

        milvus_client = AsyncMilvusClient(
            uri=settings.resolved_milvus_uri,
            token=settings.milvus_token or None,
        )

        async def embed_profile(text: str) -> list[float]:
            endpoints = await DatabaseModelEndpointResolver().resolve(
                agent_type="memory_projection", task_type="embedding"
            )
            if not endpoints:
                raise RecoverableAgentError("没有可用的 embedding 端点")
            return (await get_memory_embedding_service().embed(endpoints, text)).vector

        memory_sync_handlers["milvus"] = MilvusProfileProjection(
            milvus_client, embed_profile
        ).upsert
    else:
        logger.warning(
            "Milvus profile projection not configured; profile projection remains pending"
        )
    if neo4j_driver is not None:
        memory_sync_handlers["neo4j"] = Neo4jProfileProjection(neo4j_driver).upsert
    if memory_sync_handlers:
        memory_sync_worker = MemorySyncOutboxWorker(memory_sync_handlers)
    try:
        while True:
            try:
                worked = await runtime.run_once()
                if memory_sync_worker is not None:
                    worked = await memory_sync_worker.run_once(
                        target_store="neo4j"
                    ) or worked
                worked = await offsite_worker.run_once() or worked
            except Exception:
                # 常驻 Worker 不能因为"某一轮"的异常就整体退出：数据库抖动、
                # 迁移期间锁表、外部依赖瞬断都会命中这里，而 `run_once` 里的
                # 裸查询没有兜底。单轮失败记录堆栈后退避重试；`--once` 模式
                # 保持抛出，便于诊断一次性运行的真实问题。
                logger.warning("worker round failed; retrying after backoff", exc_info=True)
                if once:
                    raise
                await asyncio.sleep(settings.worker_poll_seconds)
                continue
            if once:
                return
            if not worked:
                await asyncio.sleep(settings.worker_poll_seconds)
    finally:
        await offsite_worker.close()
        if neo4j_driver is not None:
            await neo4j_driver.close()
        if milvus_client is not None:
            await milvus_client.close()
        await engine.dispose()


def main() -> None:
    parser = argparse.ArgumentParser(description="Agent 底座独立 Worker")
    parser.add_argument("--once", action="store_true", help="消费一次后退出")
    args = parser.parse_args()
    logging.basicConfig(level=get_settings().log_level)
    try:
        asyncio.run(serve(once=args.once))
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()