# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
# SPDX-License-Identifier: AGPL-3.0

"""Unit tests for TaskTracker."""

import json
import time

import pytest

from openviking.pyagfs.exceptions import AGFSAlreadyExistsError
from openviking.server.identity import RequestContext, Role
from openviking.service.session_service import SessionService
from openviking.service.task_store import PersistentTaskStore
from openviking.service.task_tracker import (
    TaskStatus,
    TaskTracker,
    _sanitize_error,
    get_task_tracker,
    reset_task_tracker,
    set_task_tracker,
)
from openviking_cli.session.user_id import UserIdentifier

pytestmark = pytest.mark.asyncio


@pytest.fixture(autouse=True)
def clean_singleton():
    """Reset singleton before and after each test."""
    reset_task_tracker()
    yield
    reset_task_tracker()


@pytest.fixture
def tracker() -> TaskTracker:
    return TaskTracker(store=PersistentTaskStore(_FakeAgfs()))


def _owner_kwargs(account_id: str = "acme", user_id: str = "alice"):
    return {
        "account_id": account_id,
        "user_id": user_id,
    }


def _make_ctx(account_id: str = "acme", user_id: str = "alice") -> RequestContext:
    return RequestContext(
        user=UserIdentifier(account_id, user_id),
        role=Role.ADMIN,
    )


def _set_fake_global_tracker() -> TaskTracker:
    tracker = TaskTracker(store=PersistentTaskStore(_FakeAgfs()))
    set_task_tracker(tracker)
    return tracker


class _FakeAgfs:
    def __init__(self):
        self.files = {}
        self.dirs = {"/", "/local"}

    def mkdir(self, path: str, mode: str = "755"):
        self.dirs.add(path.rstrip("/") or "/")
        return {"message": "created", "mode": mode}

    def write(self, path: str, data):
        if isinstance(data, str):
            data = data.encode("utf-8")
        self.files[path] = data
        parent = path.rsplit("/", 1)[0] or "/"
        self.dirs.add(parent)
        return "OK"

    def read(self, path: str, offset: int = 0, size: int = -1, stream: bool = False):
        if path not in self.files:
            raise FileNotFoundError(path)
        data = self.files[path]
        if size >= 0:
            return data[offset : offset + size]
        return data[offset:]

    def ls(self, path: str = "/"):
        prefix = path.rstrip("/") or "/"
        if prefix not in self.dirs:
            return []
        children = {}
        for directory in self.dirs:
            if directory in {prefix, "/"}:
                continue
            if directory.startswith(prefix + "/"):
                name = directory[len(prefix) + 1 :].split("/", 1)[0]
                if name:
                    children[name] = {"name": name, "path": f"{prefix}/{name}", "is_dir": True}
        for file_path in self.files:
            if file_path.startswith(prefix + "/"):
                name = file_path[len(prefix) + 1 :].split("/", 1)[0]
                if name and "/" not in file_path[len(prefix) + 1 :]:
                    children[name] = {"name": name, "path": f"{prefix}/{name}", "is_dir": False}
        return list(children.values())


class _FakeAgfsExistingDir(_FakeAgfs):
    def mkdir(self, path: str, mode: str = "755"):
        normalized = path.rstrip("/") or "/"
        if normalized in self.dirs:
            raise AGFSAlreadyExistsError(f"already exists: {path}")
        self.dirs.add(normalized)
        return {"message": "created", "mode": mode}


# ── Basic CRUD ──


async def test_create_task(tracker: TaskTracker):
    task = await tracker.create("session_commit", resource_id="sess-123", **_owner_kwargs())
    assert task.task_id
    assert task.task_type == "session_commit"
    assert task.resource_id == "sess-123"
    assert task.status == TaskStatus.PENDING


async def test_start_task(tracker: TaskTracker):
    task = await tracker.create("session_commit", **_owner_kwargs())
    await tracker.start(task.task_id)
    retrieved = await tracker.get(task.task_id)
    assert retrieved is not None
    assert retrieved.status == TaskStatus.RUNNING


async def test_update_stage(tracker: TaskTracker):
    task = await tracker.create("add_resource", **_owner_kwargs())
    await tracker.start(task.task_id, stage="queued")
    await tracker.update_stage(task.task_id, "parsing")
    retrieved = await tracker.get(task.task_id)
    assert retrieved is not None
    assert retrieved.status == TaskStatus.RUNNING
    assert retrieved.stage == "parsing"


async def test_complete_task(tracker: TaskTracker):
    task = await tracker.create("session_commit", resource_id="s1", **_owner_kwargs())
    await tracker.start(task.task_id)
    await tracker.complete(task.task_id, {"memories_extracted": 3})
    retrieved = await tracker.get(task.task_id)
    assert retrieved is not None
    assert retrieved.status == TaskStatus.COMPLETED
    assert retrieved.stage == "completed"
    assert retrieved.result == {"memories_extracted": 3}


async def test_task_result_redacts_user_key(tracker: TaskTracker):
    task = await tracker.create("legacy_migration", **_owner_kwargs())
    await tracker.complete(
        task.task_id,
        {
            "created_users": [{"account_id": "acme", "user_id": "bob", "user_key": "secret-key"}],
            "nested": {"user_key": "nested-secret"},
        },
    )

    retrieved = await tracker.get(task.task_id)

    assert retrieved is not None
    assert retrieved.result == {
        "created_users": [{"account_id": "acme", "user_id": "bob"}],
        "nested": {},
    }
    assert "user_key" not in json.dumps(retrieved.to_dict())


async def test_fail_task(tracker: TaskTracker):
    task = await tracker.create("session_commit", **_owner_kwargs())
    await tracker.start(task.task_id)
    await tracker.fail(task.task_id, "LLM timeout")
    retrieved = await tracker.get(task.task_id)
    assert retrieved is not None
    assert retrieved.status == TaskStatus.FAILED
    assert retrieved.stage == "failed"
    assert "LLM timeout" in retrieved.error


async def test_get_nonexistent_returns_none(tracker: TaskTracker):
    assert await tracker.get("does-not-exist") is None


# ── List / Filter ──


async def test_list_all(tracker: TaskTracker):
    await tracker.create("session_commit", resource_id="s1", **_owner_kwargs())
    await tracker.create("resource_ingest", resource_id="r1", **_owner_kwargs())
    tasks = await tracker.list_tasks()
    assert len(tasks) == 2


async def test_list_filter_by_type(tracker: TaskTracker):
    await tracker.create("session_commit", **_owner_kwargs())
    await tracker.create("resource_ingest", **_owner_kwargs())
    tasks = await tracker.list_tasks(task_type="session_commit")
    assert len(tasks) == 1
    assert tasks[0].task_type == "session_commit"


async def test_list_filter_by_status(tracker: TaskTracker):
    t1 = await tracker.create("session_commit", **_owner_kwargs())
    await tracker.create("session_commit", **_owner_kwargs())
    await tracker.start(t1.task_id)
    await tracker.complete(t1.task_id, {})

    completed = await tracker.list_tasks(status="completed")
    assert len(completed) == 1
    pending = await tracker.list_tasks(status="pending")
    assert len(pending) == 1


async def test_list_filter_by_resource_id(tracker: TaskTracker):
    await tracker.create("session_commit", resource_id="s1", **_owner_kwargs())
    await tracker.create("session_commit", resource_id="s2", **_owner_kwargs())
    tasks = await tracker.list_tasks(resource_id="s1")
    assert len(tasks) == 1
    assert tasks[0].resource_id == "s1"


async def test_get_hides_task_from_other_owner(tracker: TaskTracker):
    task = await tracker.create(
        "session_commit",
        resource_id="s1",
        account_id="acme",
        user_id="alice",
    )

    assert (
        await tracker.get(
            task.task_id,
            account_id="acme",
            user_id="bob",
        )
        is None
    )


async def test_list_tasks_filters_by_owner(tracker: TaskTracker):
    await tracker.create(
        "session_commit",
        resource_id="alice-task",
        account_id="acme",
        user_id="alice",
    )
    await tracker.create(
        "session_commit",
        resource_id="bob-task",
        account_id="acme",
        user_id="bob",
    )

    tasks = await tracker.list_tasks(account_id="acme", user_id="alice")

    assert len(tasks) == 1
    assert tasks[0].resource_id == "alice-task"


async def test_list_limit(tracker: TaskTracker):
    for i in range(10):
        await tracker.create("session_commit", resource_id=f"s{i}", **_owner_kwargs())
    tasks = await tracker.list_tasks(limit=3)
    assert len(tasks) == 3


async def test_list_order_most_recent_first(tracker: TaskTracker):
    await tracker.create("session_commit", resource_id="first", **_owner_kwargs())
    await tracker.create("session_commit", resource_id="second", **_owner_kwargs())
    tasks = await tracker.list_tasks()
    assert tasks[0].resource_id == "second"
    assert tasks[1].resource_id == "first"


# ── Duplicate detection ──


async def test_has_running_detects_pending(tracker: TaskTracker):
    await tracker.create("session_commit", resource_id="s1", **_owner_kwargs())
    assert await tracker.has_running("session_commit", "s1") is True


async def test_has_running_detects_running(tracker: TaskTracker):
    t = await tracker.create("session_commit", resource_id="s1", **_owner_kwargs())
    await tracker.start(t.task_id)
    assert await tracker.has_running("session_commit", "s1") is True


async def test_has_running_false_after_complete(tracker: TaskTracker):
    t = await tracker.create("session_commit", resource_id="s1", **_owner_kwargs())
    await tracker.start(t.task_id)
    await tracker.complete(t.task_id, {})
    assert await tracker.has_running("session_commit", "s1") is False


async def test_has_running_false_after_fail(tracker: TaskTracker):
    t = await tracker.create("session_commit", resource_id="s1", **_owner_kwargs())
    await tracker.start(t.task_id)
    await tracker.fail(t.task_id, "error")
    assert await tracker.has_running("session_commit", "s1") is False


async def test_create_if_no_running_isolated_by_owner(tracker: TaskTracker):
    alice_task = await tracker.create_if_no_running(
        "reindex",
        "viking://resources/demo",
        account_id="acme",
        user_id="alice",
    )
    bob_task = await tracker.create_if_no_running(
        "reindex",
        "viking://resources/demo",
        account_id="acme",
        user_id="bob",
    )

    assert alice_task is not None
    assert bob_task is not None
    assert alice_task.task_id != bob_task.task_id


# ── Serialization ──


async def test_to_dict(tracker: TaskTracker):
    task = await tracker.create(
        "session_commit",
        resource_id="s1",
        **_owner_kwargs(),
    )
    d = task.to_dict()
    assert d["task_id"] == task.task_id
    assert d["status"] == "pending"
    assert d["task_type"] == "session_commit"
    assert d["resource_id"] == "s1"
    assert d["stage"] is None
    assert isinstance(d["created_at"], float)
    assert isinstance(d["updated_at"], float)
    assert isinstance(d["created_at_iso"], str)
    assert "T" in d["created_at_iso"]
    assert isinstance(d["updated_at_iso"], str)
    assert "account_id" not in d
    assert "user_id" not in d


# ── Sanitization ──


async def test_sanitize_removes_sk_key():
    assert "[REDACTED]" in _sanitize_error("Error with sk-ant-api03-DAqSxxxxx")


async def test_sanitize_removes_ghp_token():
    assert "[REDACTED]" in _sanitize_error("Auth failed ghp_" + "x" * 36)


async def test_sanitize_removes_bearer_token():
    assert "[REDACTED]" in _sanitize_error("Bearer xoxb-1234567890-abcdefghij")


async def test_sanitize_truncates_long_error():
    long_error = "x" * 1000
    sanitized = _sanitize_error(long_error)
    assert len(sanitized) <= 520  # 500 + "...[truncated]"
    assert sanitized.endswith("...[truncated]")


async def test_sanitize_preserves_safe_error():
    safe = "LLM timeout after 30s"
    assert _sanitize_error(safe) == safe


# ── TTL / Eviction ──


async def test_evict_expired_completed(tracker: TaskTracker):
    t = await tracker.create("session_commit", **_owner_kwargs())
    await tracker.start(t.task_id)
    await tracker.complete(t.task_id, {})
    # Simulate old timestamp (access internal state; get() returns defensive copies)
    tracker._tasks[t.task_id].updated_at = time.time() - tracker.TTL_COMPLETED - 1
    await tracker._evict_expired()
    assert await tracker.get(t.task_id) is None


async def test_evict_keeps_recent_completed(tracker: TaskTracker):
    t = await tracker.create("session_commit", **_owner_kwargs())
    await tracker.start(t.task_id)
    await tracker.complete(t.task_id, {})
    await tracker._evict_expired()
    assert await tracker.get(t.task_id) is not None


async def test_evict_fifo_when_over_limit(tracker: TaskTracker):
    tracker.MAX_TASKS = 5
    tasks = []
    for i in range(7):
        tasks.append(await tracker.create("session_commit", resource_id=f"s{i}", **_owner_kwargs()))
    await tracker._evict_expired()
    assert tracker.count() == 5
    # Oldest should be gone
    assert await tracker.get(tasks[0].task_id) is None
    assert await tracker.get(tasks[1].task_id) is None
    # Newest should remain
    assert await tracker.get(tasks[6].task_id) is not None


# ── Singleton ──


async def test_singleton():
    t1 = _set_fake_global_tracker()
    t2 = get_task_tracker()
    assert t1 is t2


async def test_singleton_reset():
    t1 = _set_fake_global_tracker()
    reset_task_tracker()
    t2 = _set_fake_global_tracker()
    assert t1 is not t2


async def test_get_task_tracker_requires_service_initialization():
    with pytest.raises(RuntimeError, match="TaskTracker not initialized"):
        get_task_tracker()


async def test_persistent_store_cross_tracker_visibility():
    agfs = _FakeAgfs()
    store = PersistentTaskStore(agfs)
    tracker1 = TaskTracker(store=store)
    tracker2 = TaskTracker(store=store)

    task = await tracker1.create("session_commit", resource_id="sess-123", **_owner_kwargs())
    await tracker1.start(task.task_id, account_id="acme", user_id="alice")
    await tracker1.complete(task.task_id, {"ok": True}, account_id="acme", user_id="alice")

    loaded = await tracker2.get(task.task_id, account_id="acme", user_id="alice")

    assert loaded is not None
    assert loaded.status == TaskStatus.COMPLETED
    assert loaded.result == {"ok": True}


async def test_persistent_store_writes_task_record_json():
    agfs = _FakeAgfs()
    store = PersistentTaskStore(agfs)
    tracker = TaskTracker(store=store)

    task = await tracker.create(
        "add_resource",
        resource_id="viking://resources/demo",
        **_owner_kwargs(),
    )

    raw = agfs.files[f"/local/acme/_system/tasks/alice/{task.task_id}.json"]
    payload = json.loads(raw.decode("utf-8"))

    assert payload["task_id"] == task.task_id
    assert payload["task_type"] == "add_resource"
    assert payload["account_id"] == "acme"
    assert payload["user_id"] == "alice"
    assert payload["stage"] is None
    assert "schema_version" not in payload


async def test_persistent_store_keeps_tasktracker_tasks_dict():
    tracker = TaskTracker(store=PersistentTaskStore(_FakeAgfs()))
    task = await tracker.create("session_commit", **_owner_kwargs())
    assert task.task_id in tracker._tasks


async def test_persistent_store_survives_tracker_reset():
    agfs = _FakeAgfs()
    tracker1 = TaskTracker(store=PersistentTaskStore(agfs))
    task = await tracker1.create("session_commit", resource_id="sess-123", **_owner_kwargs())
    await tracker1.start(task.task_id, account_id="acme", user_id="alice")

    tracker2 = TaskTracker(store=PersistentTaskStore(agfs))
    loaded = await tracker2.get(task.task_id, account_id="acme", user_id="alice")

    assert loaded is not None
    assert loaded.status == TaskStatus.RUNNING


async def test_persistent_store_ignores_existing_task_dirs():
    agfs = _FakeAgfsExistingDir()
    tracker = TaskTracker(store=PersistentTaskStore(agfs))

    first = await tracker.create("session_commit", resource_id="sess-1", **_owner_kwargs())
    second = await tracker.create("session_commit", resource_id="sess-2", **_owner_kwargs())

    assert first.task_id != second.task_id
    assert agfs.files[f"/local/acme/_system/tasks/alice/{first.task_id}.json"]
    assert agfs.files[f"/local/acme/_system/tasks/alice/{second.task_id}.json"]


async def test_create_requires_owner(tracker: TaskTracker):
    with pytest.raises(TypeError):
        await tracker.create("session_commit", resource_id="sess-123")


async def test_create_if_no_running_requires_owner(tracker: TaskTracker):
    with pytest.raises(TypeError):
        await tracker.create_if_no_running("reindex", "viking://resources/demo")


async def test_create_rejects_blank_owner_values(tracker: TaskTracker):
    with pytest.raises(ValueError, match="Task ownership requires"):
        await tracker.create(
            "session_commit",
            resource_id="sess-123",
            account_id="",
            user_id="alice",
        )


async def test_session_service_get_commit_task_is_owner_scoped():
    tracker = _set_fake_global_tracker()
    task = await tracker.create("session_commit", resource_id="sess-123", **_owner_kwargs())
    service = SessionService()

    owner_result = await service.get_commit_task(task.task_id, _make_ctx())
    other_result = await service.get_commit_task(task.task_id, _make_ctx(user_id="bob"))

    assert owner_result is not None
    assert owner_result["task_id"] == task.task_id
    assert owner_result["resource_id"] == "sess-123"
    assert other_result is None


async def test_session_service_get_commit_task_also_filters_account():
    tracker = _set_fake_global_tracker()
    task = await tracker.create("session_commit", resource_id="sess-123", **_owner_kwargs())
    service = SessionService()

    other_account_result = await service.get_commit_task(
        task.task_id,
        _make_ctx(account_id="other-acme", user_id="alice"),
    )

    assert other_account_result is None
