# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
# SPDX-License-Identifier: AGPL-3.0
"""
Async Task Tracker for OpenViking.

Provides a lightweight registry for tracking background operations
(e.g. session commit with wait=false). Callers receive a task_id that can be
polled via the /tasks API to check completion status, results, or errors.

Design decisions:
  - Thread-safe (QueueManager workers run in separate threads).
  - TTL-based cleanup applies to the process-local cache.
  - Error messages are sanitized to avoid leaking sensitive data.
"""

import asyncio
import re
import threading
import time
from copy import deepcopy
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Dict, List, Optional
from uuid import uuid4

from openviking.service.task_store import TaskStore
from openviking_cli.utils.logger import get_logger

logger = get_logger(__name__)


class TaskStatus(str, Enum):
    """Lifecycle states of an async task."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class TaskRecord:
    """Immutable snapshot of an async task."""

    task_id: str
    task_type: str  # e.g. "session_commit"
    status: TaskStatus = TaskStatus.PENDING
    created_at: float = field(default_factory=time.time)
    updated_at: float = field(default_factory=time.time)
    resource_id: Optional[str] = None  # e.g. session_id
    account_id: Optional[str] = None
    user_id: Optional[str] = None
    stage: Optional[str] = None
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        """Serialize for JSON response."""
        d = asdict(self)
        d["status"] = self.status.value
        d["created_at_iso"] = datetime.fromtimestamp(self.created_at, tz=timezone.utc).isoformat()
        d["updated_at_iso"] = datetime.fromtimestamp(self.updated_at, tz=timezone.utc).isoformat()
        d["result"] = _sanitize_task_result(d.get("result"))
        d.pop("account_id", None)
        d.pop("user_id", None)
        return d


# ── Singleton ──

_instance: Optional["TaskTracker"] = None
_init_lock = threading.Lock()


def get_task_tracker() -> "TaskTracker":
    """Get the global TaskTracker singleton installed by service storage initialization."""
    with _init_lock:
        if _instance is None:
            logger.error(
                "TaskTracker accessed before service storage initialization; refusing to create "
                "a separate AGFS client. Ensure OpenVikingService installs the shared tracker "
                "with set_task_tracker() before task APIs are used.",
                stack_info=True,
            )
            raise RuntimeError(
                "TaskTracker not initialized. OpenVikingService must install the shared "
                "tracker with set_task_tracker() during storage initialization."
            )
        return _instance


def set_task_tracker(tracker: "TaskTracker") -> None:
    """Replace the global TaskTracker singleton."""
    global _instance
    with _init_lock:
        _instance = tracker


def reset_task_tracker() -> None:
    """Reset singleton (for testing)."""
    global _instance
    _instance = None


# ── Sanitization ──

_SENSITIVE_PATTERNS = re.compile(
    r"(sk-|cr_|ghp_|ntn_|xox[baprs]-|Bearer\s+)[a-zA-Z0-9._-]+",
    re.IGNORECASE,
)

_MAX_ERROR_LEN = 500
_SENSITIVE_RESULT_KEYS = {"user_key"}


def _sanitize_error(error: str) -> str:
    """Remove potential secrets from error messages."""
    sanitized = _SENSITIVE_PATTERNS.sub("[REDACTED]", error)
    if len(sanitized) > _MAX_ERROR_LEN:
        sanitized = sanitized[:_MAX_ERROR_LEN] + "...[truncated]"
    return sanitized


def _sanitize_task_result(result: Any) -> Any:
    """Remove sensitive fields from task results before exposing snapshots."""
    if isinstance(result, dict):
        return {
            key: _sanitize_task_result(value)
            for key, value in result.items()
            if key not in _SENSITIVE_RESULT_KEYS
        }
    if isinstance(result, list):
        return [_sanitize_task_result(item) for item in result]
    return result


# ── TaskTracker ──


class TaskTracker:
    """Async task tracker with persistent storage and a process-local cache.

    Async lifecycle operations are serialized by ``_async_lock``. The thread
    lock only protects sync snapshot reads of the local cache.
    """

    MAX_TASKS = 10_000
    TTL_COMPLETED = 86_400  # 24 hours
    TTL_FAILED = 604_800  # 7 days
    CLEANUP_INTERVAL = 300  # 5 minutes

    def __init__(self, store: TaskStore) -> None:
        self._store = store
        self._tasks: Dict[str, TaskRecord] = {}
        self._lock = threading.Lock()
        self._async_lock = asyncio.Lock()
        self._cleanup_task: Optional[asyncio.Task] = None
        logger.info(
            "[TaskTracker] Initialized (store=%s, max_tasks=%d)",
            self._store.__class__.__name__,
            self.MAX_TASKS,
        )

    # ── Lifecycle ──

    def start_cleanup_loop(self) -> None:
        """Start the background TTL cleanup coroutine.

        Safe to call multiple times; subsequent calls are no-ops.
        Must be called from within a running event loop.
        """
        if self._cleanup_task is not None and not self._cleanup_task.done():
            return
        self._cleanup_task = asyncio.create_task(self._cleanup_loop())
        logger.debug("[TaskTracker] Cleanup loop started")

    def stop_cleanup_loop(self) -> None:
        """Cancel the background cleanup task. Safe to call if not started."""
        if self._cleanup_task is not None and not self._cleanup_task.done():
            self._cleanup_task.cancel()
            logger.debug("[TaskTracker] Cleanup loop stopped")

    async def _cleanup_loop(self) -> None:
        while True:
            try:
                await asyncio.sleep(self.CLEANUP_INTERVAL)
                await self._evict_expired()
            except asyncio.CancelledError:
                break
            except Exception:
                logger.exception("[TaskTracker] Cleanup error")

    async def _evict_expired(self) -> None:
        """Remove expired tasks and enforce MAX_TASKS."""
        now = time.time()
        async with self._async_lock:
            with self._lock:
                expired_ids = []
                for tid, t in self._tasks.items():
                    if (
                        t.status == TaskStatus.COMPLETED
                        and (now - t.updated_at) > self.TTL_COMPLETED
                    ):
                        expired_ids.append(tid)
                    elif t.status == TaskStatus.FAILED and (now - t.updated_at) > self.TTL_FAILED:
                        expired_ids.append(tid)

                for tid in expired_ids:
                    self._tasks.pop(tid, None)

                if len(self._tasks) > self.MAX_TASKS:
                    sorted_tasks = sorted(self._tasks.items(), key=lambda x: x[1].created_at)
                    excess = len(self._tasks) - self.MAX_TASKS
                    for tid, _ in sorted_tasks[:excess]:
                        self._tasks.pop(tid, None)

            if expired_ids:
                logger.debug("[TaskTracker] Evicted %d expired tasks", len(expired_ids))

    @staticmethod
    def _matches_owner(
        task: TaskRecord,
        account_id: Optional[str] = None,
        user_id: Optional[str] = None,
    ) -> bool:
        """Return True when a task belongs to the requested owner filter."""
        if account_id is not None and task.account_id != account_id:
            return False
        if user_id is not None and task.user_id != user_id:
            return False
        return True

    @staticmethod
    def _validate_owner(account_id: str, user_id: str) -> None:
        """Reject ownerless task creation for user-originated background work."""
        if not account_id or not user_id:
            raise ValueError("Task ownership requires non-empty account_id and user_id")

    # ── CRUD ──

    async def create(
        self,
        task_type: str,
        resource_id: Optional[str] = None,
        *,
        account_id: str,
        user_id: str,
    ) -> TaskRecord:
        """Register a new pending task. Returns a snapshot copy."""
        self._validate_owner(account_id, user_id)
        task = TaskRecord(
            task_id=str(uuid4()),
            task_type=task_type,
            resource_id=resource_id,
            account_id=account_id,
            user_id=user_id,
        )
        async with self._async_lock:
            await self._store.create(task)
            with self._lock:
                self._tasks[task.task_id] = task
        logger.debug(
            "[TaskTracker] Created task %s type=%s resource=%s",
            task.task_id,
            task_type,
            resource_id,
        )
        return self._copy(task)

    async def create_if_no_running(
        self,
        task_type: str,
        resource_id: str,
        *,
        account_id: str,
        user_id: str,
    ) -> Optional[TaskRecord]:
        """Atomically check for running tasks and create a new one if none exist.

        Returns TaskRecord on success, None if a running task already exists.
        This eliminates the race condition between has_running() and create().
        """
        self._validate_owner(account_id, user_id)
        async with self._async_lock:
            for task in await self._load_all_from_store(account_id, user_id):
                with self._lock:
                    self._tasks[task.task_id] = task

            with self._lock:
                tasks = list(self._tasks.values())
            has_active = any(
                t.task_type == task_type
                and t.resource_id == resource_id
                and self._matches_owner(t, account_id, user_id)
                and t.status in (TaskStatus.PENDING, TaskStatus.RUNNING)
                for t in tasks
            )
            if has_active:
                return None
            task = TaskRecord(
                task_id=str(uuid4()),
                task_type=task_type,
                resource_id=resource_id,
                account_id=account_id,
                user_id=user_id,
            )
            await self._store.create(task)
            with self._lock:
                self._tasks[task.task_id] = task
        logger.debug(
            "[TaskTracker] Created task %s type=%s resource=%s",
            task.task_id,
            task_type,
            resource_id,
        )
        return self._copy(task)

    async def start(
        self,
        task_id: str,
        account_id: Optional[str] = None,
        user_id: Optional[str] = None,
        stage: Optional[str] = None,
    ) -> None:
        """Transition task to RUNNING."""
        async with self._async_lock:
            task = await self._load_for_update(task_id, account_id, user_id)
            if task:
                task.status = TaskStatus.RUNNING
                if stage is not None:
                    task.stage = stage
                task.updated_at = time.time()
                await self._store.update(task)
                with self._lock:
                    self._tasks[task.task_id] = task

    async def update_stage(
        self,
        task_id: str,
        stage: str,
        account_id: Optional[str] = None,
        user_id: Optional[str] = None,
    ) -> None:
        """Update task stage without changing its lifecycle status."""
        async with self._async_lock:
            task = await self._load_for_update(task_id, account_id, user_id)
            if task:
                task.stage = stage
                task.updated_at = time.time()
                await self._store.update(task)
                with self._lock:
                    self._tasks[task.task_id] = task

    async def complete(
        self,
        task_id: str,
        result: Optional[Dict[str, Any]] = None,
        account_id: Optional[str] = None,
        user_id: Optional[str] = None,
    ) -> None:
        """Transition task to COMPLETED with optional result."""
        async with self._async_lock:
            task = await self._load_for_update(task_id, account_id, user_id)
            if task:
                task.status = TaskStatus.COMPLETED
                task.stage = "completed"
                task.result = result
                task.updated_at = time.time()
                await self._store.update(task)
                with self._lock:
                    self._tasks[task.task_id] = task
        logger.info("[TaskTracker] Task %s completed", task_id)

    async def fail(
        self,
        task_id: str,
        error: str,
        account_id: Optional[str] = None,
        user_id: Optional[str] = None,
    ) -> None:
        """Transition task to FAILED with sanitized error."""
        async with self._async_lock:
            task = await self._load_for_update(task_id, account_id, user_id)
            if task:
                task.status = TaskStatus.FAILED
                task.stage = "failed"
                task.error = _sanitize_error(error)
                task.updated_at = time.time()
                await self._store.update(task)
                with self._lock:
                    self._tasks[task.task_id] = task
        logger.warning("[TaskTracker] Task %s failed: %s", task_id, _sanitize_error(error))

    async def get(
        self,
        task_id: str,
        account_id: Optional[str] = None,
        user_id: Optional[str] = None,
    ) -> Optional[TaskRecord]:
        """Look up a single task. Returns a snapshot copy (None if not found)."""
        async with self._async_lock:
            with self._lock:
                task = self._tasks.get(task_id)
            if task is None and account_id is not None:
                task = await self._load_from_store(task_id, account_id, user_id)
                if task is not None:
                    with self._lock:
                        self._tasks[task.task_id] = task
            if task is None or not self._matches_owner(task, account_id, user_id):
                return None
            return self._copy(task)

    async def list_tasks(
        self,
        task_type: Optional[str] = None,
        status: Optional[str] = None,
        resource_id: Optional[str] = None,
        limit: int = 50,
        account_id: Optional[str] = None,
        user_id: Optional[str] = None,
    ) -> List[TaskRecord]:
        """List tasks with optional filters. Most-recent first. Returns snapshot copies."""
        async with self._async_lock:
            if account_id is not None:
                loaded = await self._load_all_from_store(account_id, user_id)
                if loaded:
                    with self._lock:
                        for task in loaded:
                            self._tasks[task.task_id] = task
            with self._lock:
                source = list(self._tasks.values())
            tasks = [self._copy(t) for t in source if self._matches_owner(t, account_id, user_id)]
        if task_type:
            tasks = [t for t in tasks if t.task_type == task_type]
        if status:
            tasks = [t for t in tasks if t.status.value == status]
        if resource_id:
            tasks = [t for t in tasks if t.resource_id == resource_id]
        tasks.sort(key=lambda t: t.created_at, reverse=True)
        return tasks[:limit]

    async def has_running(
        self,
        task_type: str,
        resource_id: str,
        account_id: Optional[str] = None,
        user_id: Optional[str] = None,
    ) -> bool:
        """Check if there is already a running task for the given type+resource."""
        async with self._async_lock:
            if account_id is not None:
                loaded = await self._load_all_from_store(account_id, user_id)
                if loaded:
                    with self._lock:
                        for task in loaded:
                            self._tasks[task.task_id] = task
            with self._lock:
                tasks = list(self._tasks.values())
            return any(
                t.task_type == task_type
                and t.resource_id == resource_id
                and self._matches_owner(t, account_id, user_id)
                and t.status in (TaskStatus.PENDING, TaskStatus.RUNNING)
                for t in tasks
            )

    async def _load_for_update(
        self,
        task_id: str,
        account_id: Optional[str],
        user_id: Optional[str],
    ) -> Optional[TaskRecord]:
        with self._lock:
            task = self._tasks.get(task_id)
        if task is not None:
            return task if self._matches_owner(task, account_id, user_id) else None
        if account_id is None or user_id is None:
            return None
        return await self._load_from_store(task_id, account_id, user_id)

    @staticmethod
    def _record_from_payload(payload: Dict[str, Any]) -> TaskRecord:
        data = dict(payload)
        data["status"] = TaskStatus(data["status"])
        return TaskRecord(**data)

    async def _load_from_store(
        self,
        task_id: str,
        account_id: str,
        user_id: Optional[str],
    ) -> Optional[TaskRecord]:
        payload = await self._store.get(task_id, account_id=account_id, user_id=user_id)
        if payload is None:
            return None
        return self._record_from_payload(payload)

    async def _load_all_from_store(
        self, account_id: str, user_id: Optional[str]
    ) -> List[TaskRecord]:
        return [
            self._record_from_payload(payload)
            for payload in await self._store.list(account_id, user_id=user_id)
        ]

    @staticmethod
    def _copy(task: TaskRecord) -> TaskRecord:
        """Return a defensive copy of a TaskRecord."""
        copied = deepcopy(task)
        copied.result = _sanitize_task_result(copied.result)
        return copied

    def count(self) -> int:
        """Return total task count."""
        with self._lock:
            return len(self._tasks)

    def snapshot_counts_by_type(self) -> Dict[str, Dict[str, int]]:
        """Return a snapshot of task counts grouped by task_type and status."""
        from collections import defaultdict

        grouped: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int))  # type: ignore[assignment]
        with self._lock:
            tasks = list(self._tasks.values())
        for t in tasks:
            grouped[t.task_type][t.status.value] += 1
        return {k: dict(v) for k, v in grouped.items()}
