import unittest
from unittest.mock import Mock, patch

from api.tasks.memoir_tasks import generate_chapter_images


class GenerateChapterImagesTaskTest(unittest.TestCase):
    @patch("api.tasks.memoir_tasks.SessionLocal")
    @patch("api.tasks.memoir_tasks.TencentCosStorageService")
    @patch("api.tasks.memoir_tasks.LiblibImageProvider")
    @patch("api.tasks.memoir_tasks.MemoirImagePromptService")
    def test_generate_chapter_images_marks_successful_item_completed(
        self,
        prompt_service_cls,
        provider_cls,
        storage_cls,
        session_local_cls,
    ):
        chapter = type(
            "ChapterStub",
            (),
            {
                "id": "chapter-1",
                "user_id": "user-1",
                "title": "童年的夏天",
                "category": "childhood",
                "content": "那条路我一直记得。\n\n{{{{IMAGE:南方小镇的青石板路}}}}",
                "images": [
                    {
                        "index": 0,
                        "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}",
                        "description": "南方小镇的青石板路",
                        "status": "pending",
                        "url": None,
                    }
                ],
            },
        )()

        db = Mock()
        db.get.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",
            "prompt_context": "childhood: 童年的夏天",
        }
        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 = b"png-bytes"
        storage_inst = storage_cls.from_env.return_value
        storage_inst.upload_bytes.return_value = "https://cos.example.com/memoirs/u1/c1/0.png"

        generate_chapter_images.run("chapter-1")

        self.assertEqual(chapter.images[0]["status"], "completed")
        self.assertEqual(chapter.images[0]["url"], "https://cos.example.com/memoirs/u1/c1/0.png")
        self.assertEqual(chapter.images[0]["prompt"], "A serene southern China town")
        db.commit.assert_called()

    @patch("api.tasks.memoir_tasks.SessionLocal")
    @patch("api.tasks.memoir_tasks.TencentCosStorageService")
    @patch("api.tasks.memoir_tasks.LiblibImageProvider")
    @patch("api.tasks.memoir_tasks.MemoirImagePromptService")
    def test_generate_chapter_images_skips_completed_items_for_idempotency(
        self,
        prompt_service_cls,
        provider_cls,
        storage_cls,
        session_local_cls,
    ):
        chapter = type(
            "ChapterStub",
            (),
            {
                "id": "chapter-1",
                "user_id": "user-1",
                "title": "童年的夏天",
                "category": "childhood",
                "content": "那条路我一直记得。",
                "images": [
                    {
                        "index": 0,
                        "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}",
                        "description": "南方小镇的青石板路",
                        "status": "completed",
                        "url": "https://cos.example.com/already-there.png",
                    }
                ],
            },
        )()

        db = Mock()
        db.get.return_value = chapter
        session_local_cls.return_value = db

        generate_chapter_images.run("chapter-1")

        provider_cls.return_value.submit_generation.assert_not_called()
        storage_cls.from_env.return_value.upload_bytes.assert_not_called()