"""U8 — TaskIQ 异步任务集成 — 文档向量化与批量索引。

包装 TaskIQ broker，提供：
- 任务参数 schema 校验（Pydantic 模型）
- per-user 并发上限
- 任务状态机（PENDING → RUNNING → COMPLETED | FAILED | CANCELLED）
- Worker liveness & sweeper（worker_heartbeat_at 超时检测）
- 错误消息净化（剥离内部路径/连接串）
- 降级模式：broker 未配置时同步执行

Redis 隔离：使用独立 db=1（与 bus/cache 的 db=0 分离），或 key 前缀 ``taskiq:``。
任务参数在 Redis 中带 TTL（默认 1 小时）。

ponytail: TaskIQ broker 为可选依赖 — 通过 TYPE_CHECKING 导入，
未安装时模块仍可加载，TaskManager 退化为同步执行。
"""

from __future__ import annotations

import asyncio
import inspect
import logging
import re
import uuid
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Protocol

from pydantic import BaseModel, ConfigDict, Field, field_validator

from agentkit.core.protocol import TaskStatus
from agentkit.rag_platform.document_processor import (
    DEFAULT_CHUNK_OVERLAP,
    DEFAULT_CHUNK_SIZE,
    DocumentProcessor,
)
from agentkit.server.task_store import InMemoryTaskStore, TaskRecord

if TYPE_CHECKING:
    from llama_index.core.embeddings import BaseEmbedding
    from llama_index.vector_stores.postgres import PGVectorStore

    from agentkit.rag_platform.store import KBStore

logger = logging.getLogger(__name__)

# 默认配置
DEFAULT_MAX_CONCURRENT_PER_USER = 3
DEFAULT_WORKER_TTL_SECONDS = 300  # worker 心跳超时阈值
DEFAULT_ERROR_MESSAGE_MAX_LENGTH = 500
DEFAULT_TASK_TTL_SECONDS = 3600  # 任务参数 TTL

# Redis 隔离 — 独立 db=1，与 bus/cache 的 db=0 分离
TASKIQ_REDIS_DB = 1
TASKIQ_KEY_PREFIX = "taskiq:"


async def _maybe_await(result: object) -> object:
    """统一处理 sync/async 调用结果 — InMemoryTaskStore 方法为 sync，RedisTaskStore 为 async。"""
    if inspect.isawaitable(result):
        return await result
    return result


# ---------------------------------------------------------------------------
# 任务参数模型 — schema 校验
# ---------------------------------------------------------------------------


class VectorizeTaskParams(BaseModel):
    """向量化任务参数 — 单文档 parse → segment → vectorize → index。"""

    model_config = ConfigDict()

    document_id: str = Field(min_length=1, max_length=128)
    kb_id: str = Field(min_length=1, max_length=128)
    file_path: str = Field(min_length=1, max_length=1024)
    file_type: str = Field(min_length=1, max_length=64)
    chunk_size: int = Field(default=DEFAULT_CHUNK_SIZE, ge=50, le=8192)
    chunk_overlap: int = Field(default=DEFAULT_CHUNK_OVERLAP, ge=0, le=1024)
    user_id: str = Field(min_length=1, max_length=128)

    @field_validator("chunk_overlap")
    @classmethod
    def _overlap_less_than_size(cls, v: int, info: object) -> int:
        """chunk_overlap 必须小于 chunk_size。"""
        size = info.data.get("chunk_size", DEFAULT_CHUNK_SIZE)
        if v >= size:
            raise ValueError(f"chunk_overlap ({v}) must be < chunk_size ({size})")
        return v


class BatchIndexTaskParams(BaseModel):
    """批量索引任务参数 — 多文档并行向量化。"""

    model_config = ConfigDict()

    kb_id: str = Field(min_length=1, max_length=128)
    document_ids: list[str] = Field(min_length=1, max_length=100)
    user_id: str = Field(min_length=1, max_length=128)
    file_paths: dict[str, str] = Field(default_factory=dict)  # document_id -> file_path
    file_types: dict[str, str] = Field(default_factory=dict)  # document_id -> file_type
    chunk_size: int = Field(default=DEFAULT_CHUNK_SIZE, ge=50, le=8192)
    chunk_overlap: int = Field(default=DEFAULT_CHUNK_OVERLAP, ge=0, le=1024)


# ---------------------------------------------------------------------------
# 错误消息净化
# ---------------------------------------------------------------------------

# 匹配 Unix/Windows 文件路径
_PATH_RE = re.compile(
    r"(?:/[\w.\-]+)+|"  # Unix 路径 /usr/local/...
    r"(?:[A-Za-z]:\\[\w.\-]+(?:\\[\w.\-]+)*)"  # Windows 路径 C:\Users\...
)
# 匹配 Redis/PostgreSQL 连接串
_CONN_STR_RE = re.compile(
    r"(redis://|postgresql://|postgres://|mysql://)"
    r"[^\s\"']+",
    re.IGNORECASE,
)
# 匹配环境变量风格的密钥
_SECRET_RE = re.compile(
    r"(api[_-]?key|token|password|secret)[\s:=]+[\w\-]+",
    re.IGNORECASE,
)


def sanitize_error_message(msg: str, max_length: int = DEFAULT_ERROR_MESSAGE_MAX_LENGTH) -> str:
    """净化错误消息 — 剥离内部路径、连接串、密钥。

    Args:
        msg: 原始错误消息
        max_length: 最大长度（截断）

    Returns:
        净化后的错误消息
    """
    if not msg:
        return ""
    # 连接串优先于路径 — 避免 /localhost:6379/0 被路径正则部分匹配
    cleaned = _CONN_STR_RE.sub("<connection-string>", msg)
    cleaned = _PATH_RE.sub("<path>", cleaned)
    cleaned = _SECRET_RE.sub(r"\1=<redacted>", cleaned)
    if len(cleaned) > max_length:
        cleaned = cleaned[: max_length - 3] + "..."
    return cleaned


# ---------------------------------------------------------------------------
# Worker liveness & sweeper
# ---------------------------------------------------------------------------


class TaskStoreProtocol(Protocol):
    """TaskStore 协议 — InMemoryTaskStore / RedisTaskStore 均满足。"""

    async def create(
        self, task_id: str, agent_name: str, input_data: dict, skill_name: str | None = None
    ) -> TaskRecord: ...

    def get(self, task_id: str) -> TaskRecord | None: ...

    async def update_status(
        self, task_id: str, status: TaskStatus, **kwargs: object
    ) -> TaskRecord: ...

    def list_tasks(
        self, status: TaskStatus | None = None, limit: int = 100
    ) -> list[TaskRecord]: ...


class WorkerSweeper:
    """检测超时任务并标记为 failed。

    扫描 status=RUNNING 且 worker_heartbeat_at 超过 TTL 的任务，
    将其状态置为 failed，error_message="worker_timeout"。

    ponytail: 升级路径 — 当前依赖 list_tasks 全量扫描（O(n)），
    生产环境可改为 Redis ZSET 按心跳时间戳排序（O(log n)）。
    """

    def __init__(
        self,
        task_store: TaskStoreProtocol,
        ttl_seconds: int = DEFAULT_WORKER_TTL_SECONDS,
    ) -> None:
        self._store = task_store
        self._ttl = ttl_seconds

    async def sweep(self) -> int:
        """扫描超时任务，返回清理数量。"""
        now = datetime.now(timezone.utc)
        cutoff = now - timedelta(seconds=self._ttl)
        cleaned = 0

        running_tasks = await _maybe_await(
            self._store.list_tasks(status=TaskStatus.RUNNING, limit=1000)
        )

        for record in running_tasks:
            heartbeat = record.metadata.get("worker_heartbeat_at")
            if heartbeat is None:
                # 无心跳记录 — 使用 started_at 作为回退
                heartbeat = record.started_at.isoformat() if record.started_at else None
            if heartbeat is None:
                continue

            try:
                hb_time = datetime.fromisoformat(heartbeat)
            except (ValueError, TypeError):
                continue

            # 确保 timezone-aware
            if hb_time.tzinfo is None:
                hb_time = hb_time.replace(tzinfo=timezone.utc)

            if hb_time < cutoff:
                try:
                    await _maybe_await(
                        self._store.update_status(
                            record.task_id,
                            TaskStatus.FAILED,
                            error_message="worker_timeout",
                            completed_at=now,
                            metadata={**record.metadata, "swept_at": now.isoformat()},
                        )
                    )
                    cleaned += 1
                    logger.warning(
                        "Swept timed-out task %s (heartbeat=%s, cutoff=%s)",
                        record.task_id,
                        heartbeat,
                        cutoff.isoformat(),
                    )
                except KeyError:
                    # 任务已被并发删除
                    continue
        return cleaned


# ---------------------------------------------------------------------------
# TaskManager — 任务提交、查询、取消
# ---------------------------------------------------------------------------


class ConcurrencyLimitExceeded(Exception):
    """per-user 并发上限超出。"""


class TaskNotFoundError(Exception):
    """任务不存在。"""


class TaskManager:
    """异步任务管理器 — 包装 TaskIQ broker。

    如果 broker 未配置或 TaskIQ 不可用，任务同步执行（降级模式）。

    Args:
        broker: TaskIQ broker 实例（可选）
        task_store: 任务状态存储（默认 InMemoryTaskStore）
        max_concurrent_per_user: per-user 并发上限
        task_ttl_seconds: 任务参数 TTL（Redis 模式）
    """

    def __init__(
        self,
        broker: object | None = None,
        task_store: TaskStoreProtocol | None = None,
        max_concurrent_per_user: int = DEFAULT_MAX_CONCURRENT_PER_USER,
        task_ttl_seconds: int = DEFAULT_TASK_TTL_SECONDS,
    ) -> None:
        self._broker = broker
        self._store: TaskStoreProtocol = task_store or InMemoryTaskStore()
        self._max_concurrent = max_concurrent_per_user
        self._task_ttl = task_ttl_seconds
        # user_id -> 当前 RUNNING 任务数
        self._running_counts: dict[str, int] = {}
        self._lock = asyncio.Lock()

    @property
    def store(self) -> TaskStoreProtocol:
        """暴露底层 task store（供 sweeper 使用）。"""
        return self._store

    @property
    def broker(self) -> object:
        """暴露底层 broker（供 startup/shutdown 集成）。"""
        return self._broker

    # ── 任务提交 ────────────────────────────────────────────────

    async def submit_vectorize(
        self,
        params: VectorizeTaskParams,
        dependencies: dict[str, object] | None = None,
    ) -> str:
        """提交向量化任务，返回 task_id。

        Args:
            params: 向量化任务参数（已通过 schema 校验）
            dependencies: 运行时依赖（store/vector_store/embed_model），
                降级模式下直接使用；broker 模式下由 worker 注入

        Raises:
            ConcurrencyLimitExceeded: per-user 并发上限超出
        """
        await self._check_concurrency(params.user_id)

        task_id = f"vec-{uuid.uuid4()}"
        input_data = params.model_dump()
        # 依赖对象不序列化 — 仅在降级模式内存中传递
        deps = dependencies or {}

        await _maybe_await(
            self._store.create(
                task_id=task_id,
                agent_name="rag_vectorize",
                input_data=input_data,
                skill_name="vectorize",
            )
        )
        # 记录 user_id 到 metadata 便于按用户查询
        record = await _maybe_await(self._store.get(task_id))
        if record is not None:
            record.metadata["user_id"] = params.user_id

        logger.info(
            "Submitted vectorize task %s for document=%s user=%s",
            task_id,
            params.document_id,
            params.user_id,
        )

        if self._broker is not None:
            # broker 模式 — 异步派发
            await self._dispatch_to_broker("vectorize", input_data, task_id)
        else:
            # 降级模式 — 同步执行（不阻塞调用方，使用 asyncio.create_task）
            asyncio.create_task(self._run_vectorize_degraded(task_id, params, deps))

        return task_id

    async def submit_batch_index(
        self,
        params: BatchIndexTaskParams,
        dependencies: dict[str, object] | None = None,
    ) -> str:
        """提交批量索引任务，返回 task_id。"""
        await self._check_concurrency(params.user_id)

        task_id = f"batch-{uuid.uuid4()}"
        input_data = params.model_dump()
        deps = dependencies or {}

        await _maybe_await(
            self._store.create(
                task_id=task_id,
                agent_name="rag_batch_index",
                input_data=input_data,
                skill_name="batch_index",
            )
        )
        record = await _maybe_await(self._store.get(task_id))
        if record is not None:
            record.metadata["user_id"] = params.user_id

        logger.info(
            "Submitted batch index task %s for %d documents user=%s",
            task_id,
            len(params.document_ids),
            params.user_id,
        )

        if self._broker is not None:
            await self._dispatch_to_broker("batch_index", input_data, task_id)
        else:
            asyncio.create_task(self._run_batch_index_degraded(task_id, params, deps))

        return task_id

    # ── 任务查询 ────────────────────────────────────────────────

    async def get_task_status(self, task_id: str) -> TaskRecord | None:
        """查询任务状态。"""
        return await _maybe_await(self._store.get(task_id))

    async def list_tasks(
        self,
        user_id: str | None = None,
        status: TaskStatus | None = None,
        limit: int = 100,
    ) -> list[TaskRecord]:
        """查询任务历史 — 可按 user_id 和 status 过滤。"""
        tasks = await _maybe_await(self._store.list_tasks(status=status, limit=limit))
        if user_id is not None:
            tasks = [t for t in tasks if t.metadata.get("user_id") == user_id]
        return tasks

    async def cancel_task(self, task_id: str) -> bool:
        """取消任务 — 仅 PENDING/RUNNING 可取消。

        Returns:
            True 如果取消成功
        Raises:
            TaskNotFoundError: 任务不存在
        """
        record = await _maybe_await(self._store.get(task_id))
        if record is None:
            raise TaskNotFoundError(f"Task {task_id} not found")

        if record.status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED):
            return False

        await _maybe_await(
            self._store.update_status(
                task_id,
                TaskStatus.CANCELLED,
                completed_at=datetime.now(timezone.utc),
            )
        )
        # 释放并发计数
        user_id = record.metadata.get("user_id")
        if user_id and self._running_counts.get(user_id, 0) > 0:
            self._running_counts[user_id] -= 1
        return True

    # ── 心跳更新 ────────────────────────────────────────────────

    async def update_heartbeat(self, task_id: str) -> None:
        """更新 worker 心跳时间戳 — 由 worker 周期性调用。"""
        record = await _maybe_await(self._store.get(task_id))
        if record is None:
            return
        now = datetime.now(timezone.utc).isoformat()
        record.metadata["worker_heartbeat_at"] = now

    # ── 内部方法 ────────────────────────────────────────────────

    async def _check_concurrency(self, user_id: str) -> None:
        """检查 per-user 并发上限。"""
        async with self._lock:
            current = self._running_counts.get(user_id, 0)
            if current >= self._max_concurrent:
                raise ConcurrencyLimitExceeded(
                    f"User {user_id} has {current} running tasks (max={self._max_concurrent})"
                )
            self._running_counts[user_id] = current + 1

    async def _release_concurrency(self, user_id: str) -> None:
        """释放并发计数。"""
        async with self._lock:
            if self._running_counts.get(user_id, 0) > 0:
                self._running_counts[user_id] -= 1

    async def _dispatch_to_broker(self, task_type: str, params: dict, task_id: str) -> None:
        """派发任务到 TaskIQ broker。"""
        # broker.kiq 返回 TaskiqTask — 实际执行由 worker 进程完成
        try:
            await self._broker.kiq(task_type=task_type, params=params, task_id=task_id)
        except Exception as e:
            logger.error("Failed to dispatch task %s to broker: %s", task_id, e)
            raise

    async def _run_vectorize_degraded(
        self,
        task_id: str,
        params: VectorizeTaskParams,
        deps: dict[str, object],
    ) -> None:
        """降级模式 — 同步执行向量化任务（在 asyncio 任务中）。"""
        now = datetime.now(timezone.utc)
        try:
            await _maybe_await(
                self._store.update_status(task_id, TaskStatus.RUNNING, started_at=now)
            )
            await self.update_heartbeat(task_id)

            await run_vectorize_task(params, **deps)

            await _maybe_await(
                self._store.update_status(
                    task_id,
                    TaskStatus.COMPLETED,
                    completed_at=datetime.now(timezone.utc),
                    progress=1.0,
                )
            )
        except Exception as e:
            sanitized = sanitize_error_message(str(e))
            await _maybe_await(
                self._store.update_status(
                    task_id,
                    TaskStatus.FAILED,
                    error_message=sanitized,
                    completed_at=datetime.now(timezone.utc),
                )
            )
            logger.error("Vectorize task %s failed: %s", task_id, sanitized)
        finally:
            await self._release_concurrency(params.user_id)

    async def _run_batch_index_degraded(
        self,
        task_id: str,
        params: BatchIndexTaskParams,
        deps: dict[str, object],
    ) -> None:
        """降级模式 — 同步执行批量索引任务。"""
        now = datetime.now(timezone.utc)
        total = len(params.document_ids)
        succeeded: list[str] = []
        failed: list[str] = []

        try:
            await _maybe_await(
                self._store.update_status(task_id, TaskStatus.RUNNING, started_at=now)
            )
            await self.update_heartbeat(task_id)

            for i, doc_id in enumerate(params.document_ids):
                await self.update_heartbeat(task_id)
                try:
                    single_params = VectorizeTaskParams(
                        document_id=doc_id,
                        kb_id=params.kb_id,
                        file_path=params.file_paths.get(doc_id, ""),
                        file_type=params.file_types.get(doc_id, "txt"),
                        chunk_size=params.chunk_size,
                        chunk_overlap=params.chunk_overlap,
                        user_id=params.user_id,
                    )
                    await run_vectorize_task(single_params, **deps)
                    succeeded.append(doc_id)
                except Exception as e:
                    logger.warning("Batch item %s failed: %s", doc_id, e)
                    failed.append(doc_id)

                # 更新进度
                progress = (i + 1) / total
                await _maybe_await(
                    self._store.update_status(
                        task_id,
                        TaskStatus.RUNNING,
                        progress=progress,
                        progress_message=f"Processed {i + 1}/{total}",
                    )
                )

            final_status = TaskStatus.COMPLETED if not failed else TaskStatus.PARTIALLY_COMPLETED
            await _maybe_await(
                self._store.update_status(
                    task_id,
                    final_status,
                    completed_at=datetime.now(timezone.utc),
                    progress=1.0,
                    output_data={
                        "succeeded": succeeded,
                        "failed": failed,
                        "total": total,
                    },
                )
            )
        except Exception as e:
            sanitized = sanitize_error_message(str(e))
            await _maybe_await(
                self._store.update_status(
                    task_id,
                    TaskStatus.FAILED,
                    error_message=sanitized,
                    completed_at=datetime.now(timezone.utc),
                )
            )
            logger.error("Batch index task %s failed: %s", task_id, sanitized)
        finally:
            await self._release_concurrency(params.user_id)


# ---------------------------------------------------------------------------
# 任务执行函数 — 由 worker 调用
# ---------------------------------------------------------------------------


async def run_vectorize_task(
    params: VectorizeTaskParams,
    store: "KBStore",
    vector_store: "PGVectorStore",
    embed_model: "BaseEmbedding",
) -> None:
    """向量化任务 — 由 worker 执行。

    封装 DocumentProcessor.process() — parse → segment → vectorize → index。
    状态转换由 DocumentProcessor 内部管理（pending → parsing → ... → indexed | failed）。

    Args:
        params: 向量化任务参数
        store: KBStore 实例（用于文档状态更新）
        vector_store: LlamaIndex PGVectorStore 实例
        embed_model: LlamaIndex embedding 模型

    Raises:
        Exception: 管道任一阶段失败
    """
    processor = DocumentProcessor(
        chunk_size=params.chunk_size,
        chunk_overlap=params.chunk_overlap,
    )
    await processor.process(
        params.file_path,
        params.file_type,
        params.kb_id,
        params.document_id,
        vector_store,
        embed_model,
        store,
    )


async def run_batch_index_task(
    params: BatchIndexTaskParams,
    store: "KBStore",
    vector_store: "PGVectorStore",
    embed_model: "BaseEmbedding",
) -> dict[str, list[str]]:
    """批量索引任务 — 顺序处理多个文档。

    Returns:
        {"succeeded": [...], "failed": [...]}
    """
    succeeded: list[str] = []
    failed: list[str] = []

    for doc_id in params.document_ids:
        try:
            single_params = VectorizeTaskParams(
                document_id=doc_id,
                kb_id=params.kb_id,
                file_path=params.file_paths.get(doc_id, ""),
                file_type=params.file_types.get(doc_id, "txt"),
                chunk_size=params.chunk_size,
                chunk_overlap=params.chunk_overlap,
                user_id=params.user_id,
            )
            await run_vectorize_task(single_params, store, vector_store, embed_model)
            succeeded.append(doc_id)
        except Exception as e:
            logger.warning("Batch item %s failed: %s", doc_id, e)
            failed.append(doc_id)

    return {"succeeded": succeeded, "failed": failed}


# ---------------------------------------------------------------------------
# TaskIQ broker 工厂（可选 — 仅 Redis 可用时使用）
# ---------------------------------------------------------------------------


def create_broker(redis_url: str) -> object:
    """创建 TaskIQ Redis broker — 独立 db=1 隔离。

    Args:
        redis_url: Redis 连接 URL（db=0 用于 bus/cache，此处强制改为 db=1）

    Returns:
        配置好的 TaskiqBroker 实例

    Raises:
        ImportError: taskiq 未安装
    """
    from urllib.parse import urlparse, urlunparse

    from taskiq_redis import ListQueueBroker, RedisAsyncResultBackend

    from taskiq import TaskiqBroker

    # 强制使用 db=1 — 与 bus/cache 的 db=0 隔离
    parsed = urlparse(redis_url)
    # 替换 path 为 /1
    new_path = "/1"
    isolated_url = urlunparse(parsed._replace(path=new_path))

    broker = TaskiqBroker(ListQueueBroker(url=isolated_url))
    broker.with_result_backend(
        RedisAsyncResultBackend(redis_url=isolated_url, result_ex_time=DEFAULT_TASK_TTL_SECONDS)
    )

    logger.info("TaskIQ broker created with Redis db=1 isolation")
    return broker


__all__ = [
    "BatchIndexTaskParams",
    "ConcurrencyLimitExceeded",
    "DEFAULT_MAX_CONCURRENT_PER_USER",
    "DEFAULT_WORKER_TTL_SECONDS",
    "TaskManager",
    "TaskNotFoundError",
    "TaskStoreProtocol",
    "VectorizeTaskParams",
    "WorkerSweeper",
    "create_broker",
    "run_batch_index_task",
    "run_vectorize_task",
    "sanitize_error_message",
]