"""任务分发器 - 通过 Redis Queue 将任务分发给 Agent

与业务系统解耦：通过依赖注入获取 Redis 连接和数据库会话。
"""

import ipaddress
import json
import logging
import uuid
from datetime import datetime, timezone
from typing import Any, Callable, Awaitable
from urllib.parse import urlparse

from agentkit.core.exceptions import (
    NoAvailableAgentError,
    TaskDispatchError,
    TaskNotFoundError,
)
from agentkit.core.protocol import (
    AgentStatus,
    TaskMessage,
    TaskProgress,
    TaskResult,
    TaskStatus,
)

logger = logging.getLogger(__name__)

_PRIVATE_NETWORKS = [
    ipaddress.ip_network("127.0.0.0/8"),
    ipaddress.ip_network("10.0.0.0/8"),
    ipaddress.ip_network("172.16.0.0/12"),
    ipaddress.ip_network("192.168.0.0/16"),
    ipaddress.ip_network("169.254.0.0/16"),
    ipaddress.ip_network("::1/128"),
    ipaddress.ip_network("fc00::/7"),
    ipaddress.ip_network("fe80::/10"),
]


def _validate_callback_url(url: str) -> bool:
    """Validate callback URL to prevent SSRF attacks.

    Rules:
    - Only http/https protocols allowed
    - No localhost or loopback addresses
    - No private/internal IP ranges
    - No link-local addresses

    Returns True if valid, False if should be blocked.
    """
    try:
        parsed = urlparse(url)
    except Exception:
        return False

    if parsed.scheme not in ("http", "https"):
        return False

    hostname = parsed.hostname
    if not hostname:
        return False

    if hostname.lower() in ("localhost", "127.0.0.1", "::1"):
        return False

    try:
        ip = ipaddress.ip_address(hostname)
        for network in _PRIVATE_NETWORKS:
            if ip in network:
                return False
    except ValueError:
        pass

    return True


class TaskDispatcher:
    """任务分发器，通过 Redis Queue 将任务分发给 Agent"""

    def __init__(
        self,
        redis_factory: Callable[[], Awaitable[Any]],
        session_factory: Callable[[], Any],
        agent_model: Any,
        task_model: Any,
        task_log_model: Any,
    ):
        """
        Args:
            redis_factory: 返回 Redis 连接的异步工厂
            session_factory: 返回 async context manager 的工厂，用于获取数据库会话
            agent_model: Agent ORM 模型类
            task_model: AgentTask ORM 模型类
            task_log_model: AgentTaskLog ORM 模型类
        """
        self._redis_factory = redis_factory
        self._session_factory = session_factory
        self._agent_model = agent_model
        self._task_model = task_model
        self._task_log_model = task_log_model

    async def _get_redis(self):
        return await self._redis_factory()

    async def dispatch(
        self,
        task: TaskMessage,
        organization_id: str | None = None,
        created_by: str | None = None,
    ) -> str:
        """分发任务到对应 Agent 的队列，返回 task_id"""
        async with self._session_factory() as db:
            try:
                from sqlalchemy import select

                AgentModel = self._agent_model
                TaskModel = self._task_model

                # 查找目标 Agent
                stmt = select(AgentModel).where(AgentModel.name == task.agent_name)
                result = await db.execute(stmt)
                agent = result.scalar_one_or_none()

                if not agent:
                    raise TaskDispatchError(task.task_id, f"Agent '{task.agent_name}' not found")

                if agent.status != AgentStatus.ONLINE:
                    raise TaskDispatchError(
                        task.task_id,
                        f"Agent '{task.agent_name}' is not online (status={agent.status})",
                    )

                # 写入 agent_tasks 表
                task_id_uuid = uuid.UUID(task.task_id)
                agent_task = TaskModel(
                    id=task_id_uuid,
                    agent_id=agent.id,
                    task_type=task.task_type,
                    status=TaskStatus.PENDING,
                    priority=task.priority,
                    input_data=task.input_data,
                    organization_id=uuid.UUID(organization_id) if organization_id else agent.id,
                    created_by=uuid.UUID(created_by) if created_by else None,
                )
                db.add(agent_task)
                await db.commit()

                # 推送到 Redis Queue
                redis = await self._get_redis()
                queue_key = f"agent:{task.agent_name}:tasks"
                await redis.lpush(queue_key, json.dumps(task.to_dict()))

                logger.info(
                    f"Task {task.task_id} dispatched to agent '{task.agent_name}' "
                    f"(type={task.task_type}, priority={task.priority})"
                )
                return task.task_id

            except TaskDispatchError:
                raise
            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to dispatch task {task.task_id}: {e}")
                raise TaskDispatchError(task.task_id, str(e))

    async def cancel_task(self, task_id: str):
        """取消任务"""
        async with self._session_factory() as db:
            try:
                from sqlalchemy import select

                TaskModel = self._task_model
                task_uuid = uuid.UUID(task_id)
                stmt = select(TaskModel).where(TaskModel.id == task_uuid)
                result = await db.execute(stmt)
                task = result.scalar_one_or_none()

                if not task:
                    raise TaskNotFoundError(task_id)

                if task.status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED):
                    logger.warning(f"Cannot cancel task {task_id} with status '{task.status}'")
                    return

                task.status = TaskStatus.CANCELLED
                task.completed_at = datetime.now(timezone.utc)
                await db.commit()

                await self._write_log(
                    db, task_id=task_id, agent_id=str(task.agent_id),
                    log_level="info", message="Task cancelled by user",
                )
                await db.commit()

                logger.info(f"Task {task_id} cancelled")

            except TaskNotFoundError:
                raise
            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to cancel task {task_id}: {e}")
                raise

    async def get_task_status(self, task_id: str) -> dict:
        """获取任务状态"""
        async with self._session_factory() as db:
            from sqlalchemy import select

            TaskModel = self._task_model
            task_uuid = uuid.UUID(task_id)
            stmt = select(TaskModel).where(TaskModel.id == task_uuid)
            result = await db.execute(stmt)
            task = result.scalar_one_or_none()

            if not task:
                raise TaskNotFoundError(task_id)

            return self._task_to_dict(task)

    async def handle_result(self, result: TaskResult):
        """处理 Agent 返回的结果"""
        async with self._session_factory() as db:
            try:
                from sqlalchemy import select

                TaskModel = self._task_model
                task_uuid = uuid.UUID(result.task_id)
                stmt = select(TaskModel).where(TaskModel.id == task_uuid)
                db_result = await db.execute(stmt)
                task = db_result.scalar_one_or_none()

                if not task:
                    logger.error(f"Task {result.task_id} not found when handling result")
                    return

                task.status = result.status
                task.output_data = result.output_data
                task.error_message = result.error_message
                task.started_at = result.started_at
                task.completed_at = result.completed_at
                await db.commit()

                log_level = "info" if result.status == TaskStatus.COMPLETED else "error"
                log_message = (
                    f"Task {result.status}"
                    if result.status == TaskStatus.COMPLETED
                    else f"Task failed: {result.error_message}"
                )
                await self._write_log(
                    db,
                    task_id=result.task_id,
                    agent_id=str(task.agent_id),
                    log_level=log_level,
                    message=log_message,
                    extra_metadata=result.metrics,
                )
                await db.commit()

                # 触发回调
                if result.output_data and result.output_data.get("callback_url"):
                    await self._trigger_callback(result.output_data["callback_url"], result)

                logger.info(f"Task {result.task_id} result handled (status={result.status})")

            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to handle result for task {result.task_id}: {e}")

    async def handle_progress(self, progress: TaskProgress):
        """处理进度上报"""
        async with self._session_factory() as db:
            try:
                from sqlalchemy import select

                AgentModel = self._agent_model
                stmt = select(AgentModel).where(AgentModel.name == progress.agent_name)
                result = await db.execute(stmt)
                agent = result.scalar_one_or_none()

                if not agent:
                    logger.warning(f"Agent '{progress.agent_name}' not found for progress report")
                    return

                await self._write_log(
                    db,
                    task_id=progress.task_id,
                    agent_id=str(agent.id),
                    log_level="info",
                    message=f"Progress: {progress.progress:.0%} - {progress.message}",
                    extra_metadata={
                        "progress": progress.progress,
                        "updated_at": progress.updated_at.isoformat(),
                    },
                )
                await db.commit()

            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to handle progress for task {progress.task_id}: {e}")

    async def retry_failed_tasks(self, max_retries: int = 3):
        """重试失败的任务"""
        async with self._session_factory() as db:
            try:
                from sqlalchemy import select

                TaskModel = self._task_model
                AgentModel = self._agent_model
                LogModel = self._task_log_model

                stmt = select(TaskModel).where(TaskModel.status == TaskStatus.FAILED)
                result = await db.execute(stmt)
                failed_tasks = result.scalars().all()

                retried = 0
                for task in failed_tasks:
                    log_stmt = select(LogModel).where(
                        LogModel.task_id == task.id,
                        LogModel.message.like("%retry%"),
                    )
                    log_result = await db.execute(log_stmt)
                    retry_count = len(log_result.scalars().all())

                    if retry_count < max_retries:
                        task.status = TaskStatus.PENDING
                        task.error_message = None
                        task.started_at = None
                        task.completed_at = None

                        agent_stmt = select(AgentModel).where(AgentModel.id == task.agent_id)
                        agent_result = await db.execute(agent_stmt)
                        agent = agent_result.scalar_one_or_none()

                        if agent and agent.status == AgentStatus.ONLINE:
                            task_msg = TaskMessage(
                                task_id=str(task.id),
                                agent_name=agent.name,
                                task_type=task.task_type,
                                priority=task.priority,
                                input_data=task.input_data or {},
                                callback_url=None,
                                created_at=datetime.now(timezone.utc),
                            )
                            redis = await self._get_redis()
                            queue_key = f"agent:{agent.name}:tasks"
                            await redis.lpush(queue_key, json.dumps(task_msg.to_dict()))

                            await self._write_log(
                                db,
                                task_id=str(task.id),
                                agent_id=str(agent.id),
                                log_level="info",
                                message=f"Task retry attempt {retry_count + 1}/{max_retries}",
                            )
                            retried += 1

                await db.commit()
                if retried > 0:
                    logger.info(f"Retried {retried} failed tasks")

            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to retry failed tasks: {e}")

    async def _write_log(
        self,
        db: Any,
        task_id: str,
        agent_id: str,
        log_level: str,
        message: str,
        extra_metadata: dict | None = None,
    ):
        LogModel = self._task_log_model
        log_entry = LogModel(
            task_id=uuid.UUID(task_id),
            agent_id=uuid.UUID(agent_id),
            log_level=log_level,
            message=message,
            extra_metadata=extra_metadata,
        )
        db.add(log_entry)

    async def _trigger_callback(self, callback_url: str, result: TaskResult):
        if not _validate_callback_url(callback_url):
            logger.warning(f"Callback URL rejected (SSRF protection): {callback_url}")
            return

        try:
            import httpx
            async with httpx.AsyncClient(timeout=10) as client:
                await client.post(callback_url, json=result.to_dict())
                logger.info(f"Callback triggered for task {result.task_id}")
        except Exception as e:
            logger.warning(f"Callback failed for task {result.task_id}: {e}")

    def _task_to_dict(self, task: Any) -> dict:
        return {
            "id": str(task.id),
            "agent_id": str(task.agent_id),
            "task_type": task.task_type,
            "status": task.status,
            "priority": task.priority,
            "input_data": task.input_data,
            "output_data": task.output_data,
            "error_message": task.error_message,
            "created_by": str(task.created_by) if task.created_by else None,
            "organization_id": str(task.organization_id),
            "project_id": str(task.project_id) if task.project_id else None,
            "scheduled_at": task.scheduled_at.isoformat() if task.scheduled_at else None,
            "started_at": task.started_at.isoformat() if task.started_at else None,
            "completed_at": task.completed_at.isoformat() if task.completed_at else None,
            "created_at": task.created_at.isoformat() if task.created_at else None,
        }