"""ArkVideoBackend 单元测试 — mock Ark SDK。"""

import os
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from lib.video_backends.ark import ArkVideoBackend
from lib.video_backends.base import (
    VideoCapability,
    VideoGenerationRequest,
    VideoGenerationResult,
)


@pytest.fixture
def mock_ark_client():
    client = MagicMock()
    client.content_generation = MagicMock()
    client.content_generation.tasks = MagicMock()
    return client


@pytest.fixture
def backend(mock_ark_client):
    with patch("lib.video_backends.ark.create_ark_client", return_value=mock_ark_client):
        b = ArkVideoBackend(
            api_key="test-ark-key",
        )
    b._client = mock_ark_client
    return b


def _mock_httpx_stream(data: bytes = b"fake-mp4-data"):
    """Create a patched httpx mock that supports async stream context manager."""
    patcher = patch("lib.video_backends.base.httpx")
    mock_httpx = patcher.start()

    mock_stream_response = MagicMock()
    mock_stream_response.status_code = 200
    mock_stream_response.raise_for_status = MagicMock()

    async def _aiter_bytes(chunk_size=65536):
        yield data

    mock_stream_response.aiter_bytes = _aiter_bytes
    mock_stream_response.__aenter__ = AsyncMock(return_value=mock_stream_response)
    mock_stream_response.__aexit__ = AsyncMock(return_value=None)

    mock_http_client = AsyncMock()
    mock_http_client.stream = MagicMock(return_value=mock_stream_response)
    mock_http_client.__aenter__ = AsyncMock(return_value=mock_http_client)
    mock_http_client.__aexit__ = AsyncMock(return_value=None)
    mock_httpx.AsyncClient.return_value = mock_http_client

    return patcher


class TestArkProperties:
    def test_name(self, backend):
        assert backend.name == "ark"

    def test_capabilities(self, backend):
        caps = backend.capabilities
        assert VideoCapability.TEXT_TO_VIDEO in caps
        assert VideoCapability.IMAGE_TO_VIDEO in caps
        assert VideoCapability.GENERATE_AUDIO in caps
        assert VideoCapability.SEED_CONTROL in caps
        assert VideoCapability.FLEX_TIER in caps
        assert VideoCapability.NEGATIVE_PROMPT not in caps


class TestArkGenerate:
    async def test_text_to_video(self, backend, tmp_path):
        """文生视频：无 start_image。"""
        output = tmp_path / "out.mp4"

        create_result = MagicMock()
        create_result.id = "cgt-20250101-test"
        backend._client.content_generation.tasks.create = MagicMock(return_value=create_result)

        get_result = MagicMock()
        get_result.status = "succeeded"
        get_result.content = MagicMock()
        get_result.content.video_url = "https://cdn.example.com/video.mp4"
        get_result.seed = 58944
        get_result.usage = MagicMock()
        get_result.usage.completion_tokens = 246840
        backend._client.content_generation.tasks.get = MagicMock(return_value=get_result)

        patcher = _mock_httpx_stream()
        try:
            request = VideoGenerationRequest(
                prompt="a flower field",
                output_path=output,
                duration_seconds=5,
            )
            result = await backend.generate(request)
        finally:
            patcher.stop()

        assert isinstance(result, VideoGenerationResult)
        assert result.provider == "ark"
        assert result.model == "doubao-seedance-1-5-pro-251215"
        assert result.seed == 58944
        assert result.usage_tokens == 246840
        assert result.task_id == "cgt-20250101-test"

    async def test_image_to_video(self, backend, tmp_path):
        """图生视频：有 start_image，必须带 role=first_frame。"""
        output = tmp_path / "out.mp4"
        frame = tmp_path / "scene_E1S01.png"
        frame.write_bytes(b"fake-png")

        create_result = MagicMock()
        create_result.id = "cgt-i2v-test"
        backend._client.content_generation.tasks.create = MagicMock(return_value=create_result)

        get_result = MagicMock()
        get_result.status = "succeeded"
        get_result.content = MagicMock()
        get_result.content.video_url = "https://cdn.example.com/video2.mp4"
        get_result.seed = 12345
        get_result.usage = MagicMock()
        get_result.usage.completion_tokens = 200000
        backend._client.content_generation.tasks.get = MagicMock(return_value=get_result)

        patcher = _mock_httpx_stream()
        try:
            request = VideoGenerationRequest(
                prompt="girl opens eyes",
                output_path=output,
                start_image=frame,
                generate_audio=True,
            )
            result = await backend.generate(request)
        finally:
            patcher.stop()

        assert result.provider == "ark"
        create_call = backend._client.content_generation.tasks.create
        call_kwargs = create_call.call_args
        content_arg = call_kwargs.kwargs.get("content") or call_kwargs[1].get("content")
        assert len(content_arg) == 2
        assert content_arg[1]["type"] == "image_url"
        assert content_arg[1]["image_url"]["url"].startswith("data:image/")
        assert content_arg[1]["role"] == "first_frame"

    async def test_first_last_frame_role_fields(self, backend, tmp_path):
        """首尾帧：start_image/end_image 必须分别带 role=first_frame / role=last_frame，
        且 image_url 对象不再使用 position（由 role 表达位置）。"""
        output = tmp_path / "out.mp4"
        first = tmp_path / "first.png"
        first.write_bytes(b"fake-first")
        last = tmp_path / "last.png"
        last.write_bytes(b"fake-last")

        create_result = MagicMock()
        create_result.id = "cgt-fl-test"
        backend._client.content_generation.tasks.create = MagicMock(return_value=create_result)

        get_result = MagicMock()
        get_result.status = "succeeded"
        get_result.content = MagicMock()
        get_result.content.video_url = "https://cdn.example.com/video.mp4"
        get_result.seed = None
        get_result.usage = None
        backend._client.content_generation.tasks.get = MagicMock(return_value=get_result)

        patcher = _mock_httpx_stream()
        try:
            request = VideoGenerationRequest(
                prompt="morph",
                output_path=output,
                start_image=first,
                end_image=last,
            )
            await backend.generate(request)
        finally:
            patcher.stop()

        create_kwargs = backend._client.content_generation.tasks.create.call_args.kwargs
        content_arg = create_kwargs["content"]
        image_items = [c for c in content_arg if c["type"] == "image_url"]
        assert len(image_items) == 2
        assert image_items[0]["role"] == "first_frame"
        assert image_items[1]["role"] == "last_frame"
        # role 表达位置后，不应再塞 position 到 image_url
        assert "position" not in image_items[1]["image_url"]

    async def test_reference_images_role(self, backend, tmp_path):
        """参考图：每张 reference_images 必须带 role=reference_image（Ark 多图触发条件）。"""
        output = tmp_path / "out.mp4"
        ref1 = tmp_path / "ref1.jpg"
        ref1.write_bytes(b"fake-ref-1")
        ref2 = tmp_path / "ref2.jpg"
        ref2.write_bytes(b"fake-ref-2")

        create_result = MagicMock()
        create_result.id = "cgt-refs-test"
        backend._client.content_generation.tasks.create = MagicMock(return_value=create_result)

        get_result = MagicMock()
        get_result.status = "succeeded"
        get_result.content = MagicMock()
        get_result.content.video_url = "https://cdn.example.com/video.mp4"
        get_result.seed = None
        get_result.usage = None
        backend._client.content_generation.tasks.get = MagicMock(return_value=get_result)

        patcher = _mock_httpx_stream()
        try:
            request = VideoGenerationRequest(
                prompt="[图1] 与 [图2] 对话",
                output_path=output,
                reference_images=[ref1, ref2],
            )
            await backend.generate(request)
        finally:
            patcher.stop()

        create_kwargs = backend._client.content_generation.tasks.create.call_args.kwargs
        content_arg = create_kwargs["content"]
        image_items = [c for c in content_arg if c["type"] == "image_url"]
        assert len(image_items) == 2
        assert all(item["role"] == "reference_image" for item in image_items)

    async def test_failed_task_raises(self, backend, tmp_path):
        output = tmp_path / "out.mp4"

        create_result = MagicMock()
        create_result.id = "cgt-fail"
        backend._client.content_generation.tasks.create = MagicMock(return_value=create_result)

        get_result = MagicMock()
        get_result.status = "failed"
        get_result.error = "content violation"
        backend._client.content_generation.tasks.get = MagicMock(return_value=get_result)

        request = VideoGenerationRequest(prompt="test", output_path=output)
        with pytest.raises(RuntimeError, match="Ark 视频生成失败"):
            await backend.generate(request)

    async def test_with_seed_and_flex(self, backend, tmp_path):
        output = tmp_path / "out.mp4"

        create_result = MagicMock()
        create_result.id = "cgt-flex"
        backend._client.content_generation.tasks.create = MagicMock(return_value=create_result)

        get_result = MagicMock()
        get_result.status = "succeeded"
        get_result.content = MagicMock()
        get_result.content.video_url = "https://cdn.example.com/video.mp4"
        get_result.seed = 42
        get_result.usage = MagicMock()
        get_result.usage.completion_tokens = 100000
        backend._client.content_generation.tasks.get = MagicMock(return_value=get_result)

        patcher = _mock_httpx_stream()
        try:
            request = VideoGenerationRequest(
                prompt="test",
                output_path=output,
                seed=42,
                service_tier="flex",
            )
            await backend.generate(request)
        finally:
            patcher.stop()

        create_call = backend._client.content_generation.tasks.create
        call_kwargs = create_call.call_args
        assert call_kwargs.kwargs.get("seed") == 42 or call_kwargs[1].get("seed") == 42
        assert call_kwargs.kwargs.get("service_tier") == "flex" or call_kwargs[1].get("service_tier") == "flex"

    def test_missing_api_key_raises(self):
        with patch.dict(os.environ, {}, clear=True):
            with pytest.raises(ValueError, match="Ark API Key"):
                ArkVideoBackend(api_key=None)


class TestArkRetryBehavior:
    """测试任务创建与轮询的重试分离行为。"""

    async def test_poll_transient_error_retries_without_recreating_task(self, backend, tmp_path):
        """轮询阶段瞬态错误应重试轮询，而不是重新创建任务。"""
        output = tmp_path / "out.mp4"

        create_result = MagicMock()
        create_result.id = "cgt-retry-test"
        backend._client.content_generation.tasks.create = MagicMock(return_value=create_result)

        get_success = MagicMock()
        get_success.status = "succeeded"
        get_success.content = MagicMock()
        get_success.content.video_url = "https://cdn.example.com/video.mp4"
        get_success.seed = None
        get_success.usage = None

        # 第一次轮询抛 ConnectionError，第二次成功
        backend._client.content_generation.tasks.get = MagicMock(
            side_effect=[ConnectionError("connection reset"), get_success]
        )

        patcher = _mock_httpx_stream()
        try:
            request = VideoGenerationRequest(prompt="test", output_path=output)
            with patch("lib.video_backends.base.asyncio.sleep", new_callable=AsyncMock):
                result = await backend.generate(request)
        finally:
            patcher.stop()

        assert result.task_id == "cgt-retry-test"
        # 关键断言：任务只创建了一次
        assert backend._client.content_generation.tasks.create.call_count == 1
        # 轮询调用了两次（一次失败 + 一次成功）
        assert backend._client.content_generation.tasks.get.call_count == 2

    async def test_create_retries_on_transient_error(self, backend, tmp_path):
        """任务创建阶段的瞬态错误应由 @with_retry_async 重试。"""
        output = tmp_path / "out.mp4"

        create_result = MagicMock()
        create_result.id = "cgt-create-retry"
        # 第一次创建抛 ConnectionError，第二次成功
        backend._client.content_generation.tasks.create = MagicMock(
            side_effect=[ConnectionError("connection reset"), create_result]
        )

        get_result = MagicMock()
        get_result.status = "succeeded"
        get_result.content = MagicMock()
        get_result.content.video_url = "https://cdn.example.com/video.mp4"
        get_result.seed = None
        get_result.usage = None
        backend._client.content_generation.tasks.get = MagicMock(return_value=get_result)

        patcher = _mock_httpx_stream()
        try:
            request = VideoGenerationRequest(prompt="test", output_path=output)
            with (
                patch("lib.video_backends.base.asyncio.sleep", new_callable=AsyncMock),
                patch("lib.retry.asyncio.sleep", new_callable=AsyncMock),
            ):
                result = await backend.generate(request)
        finally:
            patcher.stop()

        assert result.task_id == "cgt-create-retry"
        # 创建调用了两次（一次失败 + 一次成功）
        assert backend._client.content_generation.tasks.create.call_count == 2

    async def test_poll_non_retryable_error_propagates(self, backend, tmp_path):
        """轮询阶段不可重试的错误应直接抛出。"""
        output = tmp_path / "out.mp4"

        create_result = MagicMock()
        create_result.id = "cgt-no-retry"
        backend._client.content_generation.tasks.create = MagicMock(return_value=create_result)

        backend._client.content_generation.tasks.get = MagicMock(side_effect=ValueError("invalid response"))

        request = VideoGenerationRequest(prompt="test", output_path=output)
        with pytest.raises(ValueError, match="invalid response"):
            with patch("lib.video_backends.base.asyncio.sleep", new_callable=AsyncMock):
                await backend.generate(request)

        # 创建只调用一次，轮询只尝试一次就抛出
        assert backend._client.content_generation.tasks.create.call_count == 1
        assert backend._client.content_generation.tasks.get.call_count == 1


class TestArkModelCapabilities:
    """测试不同模型的能力映射。"""

    def test_seedance_2_no_flex_tier(self):
        with patch("lib.video_backends.ark.create_ark_client", return_value=MagicMock()):
            b = ArkVideoBackend(api_key="test", model="doubao-seedance-2-0-260128")
        caps = b.capabilities
        assert VideoCapability.FLEX_TIER not in caps
        assert VideoCapability.VIDEO_EXTEND not in caps

    def test_seedance_2_fast_no_flex_tier(self):
        with patch("lib.video_backends.ark.create_ark_client", return_value=MagicMock()):
            b = ArkVideoBackend(api_key="test", model="doubao-seedance-2-0-fast-260128")
        caps = b.capabilities
        assert VideoCapability.FLEX_TIER not in caps
        assert VideoCapability.VIDEO_EXTEND not in caps

    def test_seedance_1_5_has_flex_tier(self):
        with patch("lib.video_backends.ark.create_ark_client", return_value=MagicMock()):
            b = ArkVideoBackend(api_key="test", model="doubao-seedance-1-5-pro-251215")
        caps = b.capabilities
        assert VideoCapability.FLEX_TIER in caps
        assert VideoCapability.VIDEO_EXTEND not in caps

    def test_unknown_model_gets_default_capabilities(self):
        with patch("lib.video_backends.ark.create_ark_client", return_value=MagicMock()):
            b = ArkVideoBackend(api_key="test", model="some-future-model")
        caps = b.capabilities
        assert VideoCapability.FLEX_TIER in caps

    def test_seedance_2_dot_format_no_flex_tier(self):
        """ark-agent-plan 用 dot 命名（doubao-seedance-2.0），同样不该带 FLEX_TIER。"""
        with patch("lib.video_backends.ark.create_ark_client", return_value=MagicMock()):
            b = ArkVideoBackend(api_key="test", model="doubao-seedance-2.0")
        assert VideoCapability.FLEX_TIER not in b.capabilities

    def test_seedance_2_fast_dot_format_no_flex_tier(self):
        with patch("lib.video_backends.ark.create_ark_client", return_value=MagicMock()):
            b = ArkVideoBackend(api_key="test", model="doubao-seedance-2.0-fast")
        assert VideoCapability.FLEX_TIER not in b.capabilities


class TestArkServiceTierParam:
    """service_tier 只对声明了 FLEX_TIER 能力的模型传入，否则 API 会报错。"""

    async def test_seedance_2_does_not_send_service_tier(self, tmp_path):
        output = tmp_path / "out.mp4"
        mock_client = MagicMock()
        mock_client.content_generation = MagicMock()
        mock_client.content_generation.tasks = MagicMock()

        with patch("lib.video_backends.ark.create_ark_client", return_value=mock_client):
            backend = ArkVideoBackend(api_key="test", model="doubao-seedance-2-0-260128")
        backend._client = mock_client

        create_result = MagicMock()
        create_result.id = "cgt-seedance2"
        backend._client.content_generation.tasks.create = MagicMock(return_value=create_result)

        get_result = MagicMock()
        get_result.status = "succeeded"
        get_result.content = MagicMock()
        get_result.content.video_url = "https://cdn.example.com/v.mp4"
        get_result.seed = None
        get_result.usage = None
        backend._client.content_generation.tasks.get = MagicMock(return_value=get_result)

        patcher = _mock_httpx_stream()
        try:
            request = VideoGenerationRequest(prompt="test", output_path=output)
            with patch("lib.video_backends.base.asyncio.sleep", new_callable=AsyncMock):
                await backend.generate(request)
        finally:
            patcher.stop()

        create_kwargs = backend._client.content_generation.tasks.create.call_args.kwargs
        assert "service_tier" not in create_kwargs

    async def test_seedance_1_5_sends_service_tier(self, backend, tmp_path):
        output = tmp_path / "out.mp4"

        create_result = MagicMock()
        create_result.id = "cgt-seedance15"
        backend._client.content_generation.tasks.create = MagicMock(return_value=create_result)

        get_result = MagicMock()
        get_result.status = "succeeded"
        get_result.content = MagicMock()
        get_result.content.video_url = "https://cdn.example.com/v.mp4"
        get_result.seed = None
        get_result.usage = None
        backend._client.content_generation.tasks.get = MagicMock(return_value=get_result)

        patcher = _mock_httpx_stream()
        try:
            request = VideoGenerationRequest(prompt="test", output_path=output)
            with patch("lib.video_backends.base.asyncio.sleep", new_callable=AsyncMock):
                await backend.generate(request)
        finally:
            patcher.stop()

        create_kwargs = backend._client.content_generation.tasks.create.call_args.kwargs
        assert create_kwargs.get("service_tier") == "default"


class TestArkVideoBackendBaseUrl:
    def test_custom_base_url_passed_through(self):
        with patch("lib.video_backends.ark.create_ark_client") as mock_create:
            ArkVideoBackend(api_key="k", base_url="https://ark.cn-beijing.volces.com/api/plan/v3")
            mock_create.assert_called_once_with(
                api_key="k",
                base_url="https://ark.cn-beijing.volces.com/api/plan/v3",
            )

    def test_default_base_url_is_none(self):
        with patch("lib.video_backends.ark.create_ark_client") as mock_create:
            ArkVideoBackend(api_key="k")
            mock_create.assert_called_once_with(api_key="k", base_url=None)
