"""CitationDetector Agent - 将现有 CitationEngine 封装为 Agent"""

import logging
import re
import time
from datetime import datetime, timezone

from app.agent_framework.base import BaseAgent
from app.agent_framework.protocol import (
    AgentCapability,
    AgentType,
    TaskMessage,
    TaskResult,
    TaskStatus,
)
from app.database import AsyncSessionLocal
from app.models.citation_record import CitationRecord
from app.models.query import Query
from app.workers.citation_engine import CitationEngine

logger = logging.getLogger(__name__)


class CitationDetectorAgent(BaseAgent):
    """
    引用检测 Agent：将现有 CitationEngine 封装为 BaseAgent 实现。

    支持的任务类型:
    - citation_detect: 执行完整的引用检测（遍历 query 的所有平台）
    - citation_detect_single: 执行单个平台的引用检测
    """

    def __init__(self):
        super().__init__(
            name="citation_detector",
            agent_type=AgentType.CITATION_DETECTOR,
            version="1.0.0",
        )
        self._engine = CitationEngine()

    def get_capabilities(self) -> AgentCapability:
        return AgentCapability(
            agent_name=self.name,
            agent_type=self.agent_type,
            version=self.version,
            supported_tasks=["citation_detect", "citation_detect_single"],
            max_concurrency=3,
            description="AI平台引用检测Agent：检测目标品牌在各AI平台回答中的引用情况",
        )

    async def execute(self, task: TaskMessage) -> TaskResult:
        """执行引用检测任务"""
        started_at = datetime.now(timezone.utc)
        start_time = time.monotonic()

        try:
            if task.task_type == "citation_detect":
                output = await self._execute_full_detect(task)
            elif task.task_type == "citation_detect_single":
                output = await self._execute_single_detect(task)
            else:
                raise ValueError(f"Unsupported task type: {task.task_type}")

            elapsed = time.monotonic() - start_time
            return TaskResult(
                task_id=task.task_id,
                agent_name=self.name,
                status=TaskStatus.COMPLETED,
                output_data=output,
                error_message=None,
                started_at=started_at,
                completed_at=datetime.now(timezone.utc),
                metrics={
                    "elapsed_seconds": round(elapsed, 2),
                    "task_type": task.task_type,
                },
            )

        except Exception as e:
            elapsed = time.monotonic() - start_time
            logger.error(f"CitationDetector task {task.task_id} failed: {e}")
            return TaskResult(
                task_id=task.task_id,
                agent_name=self.name,
                status=TaskStatus.FAILED,
                output_data=None,
                error_message=str(e),
                started_at=started_at,
                completed_at=datetime.now(timezone.utc),
                metrics={
                    "elapsed_seconds": round(elapsed, 2),
                    "task_type": task.task_type,
                },
            )

    async def _execute_full_detect(self, task: TaskMessage) -> dict:
        """
        执行完整的引用检测（遍历 query 的所有平台）。
        input_data 需包含: query_id (str)
        """
        query_id = task.input_data.get("query_id")
        if not query_id:
            raise ValueError("input_data must contain 'query_id'")

        async with AsyncSessionLocal() as db:
            from sqlalchemy import select

            stmt = select(Query).where(Query.id == query_id)
            result = await db.execute(stmt)
            query = result.scalar_one_or_none()

            if not query:
                raise ValueError(f"Query {query_id} not found")

            # 上报进度：开始检测
            await self.report_progress(
                task_id=task.task_id,
                progress=0.1,
                message=f"Starting citation detection for query '{query.keyword}'",
            )

            records = await self._engine.execute_query(query, db)

            # 上报进度：检测完成
            await self.report_progress(
                task_id=task.task_id,
                progress=1.0,
                message=f"Detection completed: {len(records)} records found",
            )

            # 构建输出
            record_summaries = []
            for r in records:
                record_summaries.append({
                    "id": str(r.id),
                    "platform": r.platform,
                    "cited": r.cited,
                    "confidence": r.confidence,
                    "match_type": r.match_type,
                })

            return {
                "query_id": str(query_id),
                "keyword": query.keyword,
                "total_records": len(records),
                "cited_count": sum(1 for r in records if r.cited),
                "records": record_summaries,
            }

    async def _execute_single_detect(self, task: TaskMessage) -> dict:
        """
        执行单个平台的引用检测。
        input_data 需包含: keyword, platform, target_brand, brand_aliases(optional)
        """
        keyword = task.input_data.get("keyword")
        platform = task.input_data.get("platform")
        target_brand = task.input_data.get("target_brand")
        brand_aliases = task.input_data.get("brand_aliases", [])

        if not all([keyword, platform, target_brand]):
            raise ValueError(
                "input_data must contain 'keyword', 'platform', 'target_brand'"
            )

        # 上报进度
        await self.report_progress(
            task_id=task.task_id,
            progress=0.2,
            message=f"Querying platform '{platform}' with keyword '{keyword}'",
        )

        result = await self._engine.execute_single_platform(
            keyword=keyword,
            platform=platform,
            target_brand=target_brand,
            brand_aliases=brand_aliases,
        )

        # 上报进度：完成
        await self.report_progress(
            task_id=task.task_id,
            progress=1.0,
            message=f"Single platform detection completed on '{platform}'",
        )

        # 清理 raw_response 以避免返回过大的数据
        output = {k: v for k, v in result.items() if k != "raw_response"}
        return output

    # -----------------------------------------------------------------------
    # 向后兼容：保留原有 CitationEngine 的同步调用接口
    # -----------------------------------------------------------------------

    async def execute_query_compat(self, query: Query, db) -> list[CitationRecord]:
        """
        向后兼容方法：供现有 scheduler 继续调用。
        签名与 CitationEngine.execute_query 完全一致。
        """
        return await self._engine.execute_query(query, db)

    async def execute_single_platform_compat(
        self, keyword: str, platform: str, target_brand: str, brand_aliases: list
    ) -> dict:
        """
        向后兼容方法：供现有 scheduler 继续调用。
        签名与 CitationEngine.execute_single_platform 完全一致。
        """
        return await self._engine.execute_single_platform(
            keyword=keyword,
            platform=platform,
            target_brand=target_brand,
            brand_aliases=brand_aliases,
        )

    async def close(self):
        """关闭底层引擎"""
        await self._engine.close()
