import base64
import unittest
from types import SimpleNamespace
from unittest.mock import Mock, patch

from app.ports.image_gen import ImageResult, TaskStatus
from app.tasks.memoir_tasks import build_cos_key, generate_chapter_images


_ONE_BY_ONE_PNG = base64.b64decode(
    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+aF9sAAAAASUVORK5CYII="
)


def _cover_image_record(img_dict):
    d = dict(img_dict or {})
    return SimpleNamespace(
        id="cover-img-1",
        order_index=d.get("index", 0),
        placeholder=d.get("placeholder"),
        description=d.get("description"),
        status=d.get("status"),
        prompt=d.get("prompt"),
        url=d.get("url"),
        storage_key=d.get("storage_key"),
        provider=d.get("provider"),
        style=d.get("style"),
        size=d.get("size"),
        error=d.get("error"),
        retryable=d.get("retryable"),
        created_at=d.get("created_at"),
        updated_at=d.get("updated_at"),
    )


def _chapter_stub():
    rec = _cover_image_record(
        {
            "index": 0,
            "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}",
            "description": "南方小镇的青石板路",
            "status": "pending",
            "url": None,
        }
    )
    return SimpleNamespace(
        id="chapter-1",
        user_id="user-1",
        title="童年的夏天",
        category="childhood",
        canonical_markdown="# 标题\n\n那条路我一直记得。",
        sections=[],
        images=[rec],
        cover_image=None,
    )


class GenerateChapterImagesPersistenceTest(unittest.TestCase):
    @patch("app.tasks.memoir_tasks.get_sync_db")
    @patch("app.tasks.memoir_tasks.TencentCosStorageService")
    @patch("app.tasks.memoir_tasks.get_image_generator")
    @patch("app.tasks.memoir_tasks.ImagePromptOrchestrator")
    @patch("app.tasks.memoir_tasks._release_chapter_image_lock")
    @patch("app.tasks.memoir_tasks._acquire_chapter_image_lock", return_value=True)
    def test_successful_generation_persists_completed_status(
        self,
        _acquire_lock_mock,
        _release_lock_mock,
        prompt_service_cls,
        get_image_generator_mock,
        storage_cls,
        get_sync_db_mock,
    ):
        chapter = _chapter_stub()
        cover = chapter.images[0]
        db = Mock()
        db.execute.return_value.unique.return_value.scalar_one_or_none.return_value = (
            chapter
        )
        get_sync_db_mock.return_value.__enter__.return_value = db
        get_sync_db_mock.return_value.__exit__.return_value = False

        prompt_data = {
            "prompt": "A serene southern China town",
            "style": "watercolor",
            "size": "1024x1024",
            "prompt_context": "childhood: 童年的夏天",
        }
        prompt_service_cls.return_value.build_cover_prompt.return_value = prompt_data
        mock_gen = Mock()
        mock_gen.generate.return_value = ImageResult(
            status=TaskStatus.COMPLETED,
            task_id="",
            image_url="https://provider.example.com/1.png",
        )
        mock_gen.download_image.return_value = _ONE_BY_ONE_PNG
        get_image_generator_mock.return_value = mock_gen
        storage_cls.from_env.return_value.upload_bytes.return_value = (
            "https://cos.example.com/memoirs/user-1/chapter-1/cover.png"
        )

        generate_chapter_images.run("chapter-1")

        self.assertEqual(cover.status, "completed")
        self.assertEqual(
            cover.url,
            "https://cos.example.com/memoirs/user-1/chapter-1/cover.png",
        )
        self.assertEqual(cover.prompt, "A serene southern China town")
        self.assertEqual(
            cover.storage_key,
            build_cos_key("user-1", "chapter-1", "cover", prompt_data["prompt"]),
        )