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

from tasks.memoir_tasks import generate_chapter_images


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


def _image_record(img_dict):
    d = dict(img_dict or {})
    return SimpleNamespace(
        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 = _image_record(
        {
            "index": 0,
            "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}",
            "description": "南方小镇的青石板路",
            "status": "pending",
            "url": None,
        }
    )
    section = SimpleNamespace(
        content="那条路我一直记得。",
        image_id="image-1",
        image_record=rec,
        order_index=0,
    )
    return SimpleNamespace(
        id="chapter-1",
        user_id="user-1",
        title="童年的夏天",
        category="childhood",
        sections=[section],
        images=[],
        cover_image=None,
    )


class GenerateChapterImagesPersistenceTest(unittest.TestCase):
    @patch("tasks.memoir_tasks.SessionLocal")
    @patch("tasks.memoir_tasks.TencentCosStorageService")
    @patch("tasks.memoir_tasks.LiblibImageProvider")
    @patch("tasks.memoir_tasks.MemoirImagePromptService")
    @patch("tasks.memoir_tasks._release_chapter_image_lock")
    @patch("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,
        provider_cls,
        storage_cls,
        session_local_cls,
    ):
        chapter = _chapter_stub()
        db = Mock()
        db.execute.return_value.unique.return_value.scalar_one_or_none.return_value = chapter
        session_local_cls.return_value = db

        prompt_service_cls.return_value.build_prompt.return_value = {
            "prompt": "A serene southern China town",
            "style": "watercolor",
            "size": "1024x1024",
        }
        provider_inst = provider_cls.return_value
        provider_inst.submit_generation.return_value = {
            "status": "completed",
            "image_url": "https://provider.example.com/1.png",
        }
        provider_inst.download_image.return_value = _ONE_BY_ONE_PNG
        storage_cls.from_env.return_value.upload_bytes.return_value = (
            "https://cos.example.com/memoirs/user-1/chapter-1/0.png"
        )

        generate_chapter_images.run("chapter-1")

        record = chapter.sections[0].image_record
        self.assertEqual(record.status, "completed")
        self.assertEqual(record.url, "https://cos.example.com/memoirs/user-1/chapter-1/0.png")
        self.assertEqual(record.prompt, "A serene southern China town")