database.py

import sqlite3
import json
from contextlib import contextmanager
from config import Config

WORDS_SEED = [
    # 동물
    {"korean": "사자", "english": "lion", "emoji": "🦁", "category": "동물"},
    {"korean": "토끼", "english": "rabbit", "emoji": "🐰", "category": "동물"},
    {"korean": "코끼리", "english": "elephant", "emoji": "🐘", "category": "동물"},
    {"korean": "고양이", "english": "cat", "emoji": "🐱", "category": "동물"},
    {"korean": "강아지", "english": "dog", "emoji": "🐶", "category": "동물"},
    {"korean": "곰", "english": "bear", "emoji": "🐻", "category": "동물"},
    {"korean": "나비", "english": "butterfly", "emoji": "🦋", "category": "동물"},
    {"korean": "물고기", "english": "fish", "emoji": "🐟", "category": "동물"},
    {"korean": "펭귄", "english": "penguin", "emoji": "🐧", "category": "동물"},
    {"korean": "원숭이", "english": "monkey", "emoji": "🐒", "category": "동물"},
    {"korean": "돼지", "english": "pig", "emoji": "🐷", "category": "동물"},
    {"korean": "소", "english": "cow", "emoji": "🐮", "category": "동물"},
    {"korean": "말", "english": "horse", "emoji": "🐴", "category": "동물"},
    {"korean": "오리", "english": "duck", "emoji": "🦆", "category": "동물"},
    {"korean": "거북이", "english": "turtle", "emoji": "🐢", "category": "동물"},
    {"korean": "새", "english": "bird", "emoji": "🐦", "category": "동물"},
    {"korean": "개구리", "english": "frog", "emoji": "🐸", "category": "동물"},
    {"korean": "호랑이", "english": "tiger", "emoji": "🐯", "category": "동물"},
    {"korean": "기린", "english": "giraffe", "emoji": "🦒", "category": "동물"},
    {"korean": "여우", "english": "fox", "emoji": "🦊", "category": "동물"},
    # 자연
    {"korean": "나무", "english": "tree", "emoji": "🌳", "category": "자연"},
    {"korean": "꽃", "english": "flower", "emoji": "🌸", "category": "자연"},
    {"korean": "무지개", "english": "rainbow", "emoji": "🌈", "category": "자연"},
    {"korean": "별", "english": "star", "emoji": "⭐", "category": "자연"},
    {"korean": "달", "english": "moon", "emoji": "🌙", "category": "자연"},
    {"korean": "태양", "english": "sun", "emoji": "☀️", "category": "자연"},
    {"korean": "구름", "english": "cloud", "emoji": "☁️", "category": "자연"},
    {"korean": "산", "english": "mountain", "emoji": "⛰️", "category": "자연"},
    {"korean": "바다", "english": "ocean", "emoji": "🌊", "category": "자연"},
    {"korean": "눈사람", "english": "snowman", "emoji": "⛄", "category": "자연"},
    {"korean": "번개", "english": "lightning", "emoji": "⚡", "category": "자연"},
    {"korean": "비", "english": "rain", "emoji": "🌧️", "category": "자연"},
    # 사물/판타지
    {"korean": "집", "english": "house", "emoji": "🏠", "category": "사물"},
    {"korean": "성", "english": "castle", "emoji": "🏰", "category": "사물"},
    {"korean": "자동차", "english": "car", "emoji": "🚗", "category": "사물"},
    {"korean": "기차", "english": "train", "emoji": "🚂", "category": "사물"},
    {"korean": "풍선", "english": "balloon", "emoji": "🎈", "category": "사물"},
    {"korean": "케이크", "english": "cake", "emoji": "🎂", "category": "사물"},
    {"korean": "왕관", "english": "crown", "emoji": "👑", "category": "사물"},
    {"korean": "로봇", "english": "robot", "emoji": "🤖", "category": "사물"},
    {"korean": "자전거", "english": "bicycle", "emoji": "🚲", "category": "사물"},
    {"korean": "비행기", "english": "airplane", "emoji": "✈️", "category": "사물"},
    {"korean": "배", "english": "boat", "emoji": "⛵", "category": "사물"},
    {"korean": "우주선", "english": "spaceship", "emoji": "🚀", "category": "사물"},
    {"korean": "외계인", "english": "alien", "emoji": "👽", "category": "사물"},
    {"korean": "안경", "english": "glasses", "emoji": "👓", "category": "사물"},
    {"korean": "모자", "english": "hat", "emoji": "🎩", "category": "사물"},
    # 음식
    {"korean": "사과", "english": "apple", "emoji": "🍎", "category": "음식"},
    {"korean": "딸기", "english": "strawberry", "emoji": "🍓", "category": "음식"},
    {"korean": "수박", "english": "watermelon", "emoji": "🍉", "category": "음식"},
    {"korean": "바나나", "english": "banana", "emoji": "🍌", "category": "음식"},
    {"korean": "포도", "english": "grapes", "emoji": "🍇", "category": "음식"},
    {"korean": "당근", "english": "carrot", "emoji": "🥕", "category": "음식"},
    {"korean": "아이스크림", "english": "ice cream", "emoji": "🍦", "category": "음식"},
    {"korean": "피자", "english": "pizza", "emoji": "🍕", "category": "음식"},
    {"korean": "햄버거", "english": "hamburger", "emoji": "🍔", "category": "음식"},
    {"korean": "도넛", "english": "donut", "emoji": "🍩", "category": "음식"},
    {"korean": "사탕", "english": "candy", "emoji": "🍬", "category": "음식"},
    {"korean": "초콜릿", "english": "chocolate", "emoji": "🍫", "category": "음식"},
    {"korean": "빵", "english": "bread", "emoji": "🍞", "category": "음식"},
    {"korean": "치즈", "english": "cheese", "emoji": "🧀", "category": "음식"},
]


@contextmanager
def get_db():
    conn = sqlite3.connect(Config.DATABASE_PATH)
    conn.row_factory = sqlite3.Row
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


def init_db():
    with get_db() as conn:
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS words (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                korean TEXT NOT NULL,
                english TEXT NOT NULL,
                emoji TEXT NOT NULL,
                category TEXT DEFAULT 'general',
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP
            );
            CREATE TABLE IF NOT EXISTS drawings (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                word_id INTEGER NOT NULL REFERENCES words(id),
                image_data TEXT NOT NULL,
                file_path TEXT,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP
            );
            CREATE TABLE IF NOT EXISTS stories (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                title TEXT NOT NULL,
                content TEXT NOT NULL,
                moral TEXT,
                keywords TEXT NOT NULL,
                drawing_ids TEXT NOT NULL,
                illustration_path TEXT,
                audio_path TEXT,
                scene_data TEXT,
                layout_data TEXT,
                is_read INTEGER DEFAULT 0,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP
            );
        """)
        existing_words = {
            row["korean"] for row in conn.execute("SELECT korean FROM words")
        }
        for word in WORDS_SEED:
            if word["korean"] not in existing_words:
                conn.execute(
                    "INSERT INTO words (korean, english, emoji, category) VALUES (?, ?, ?, ?)",
                    (word["korean"], word["english"], word["emoji"], word["category"]),
                )
    # 기존 DB 마이그레이션: scene_data 및 is_read 컬럼 추가
    try:
        with get_db() as conn:
            conn.execute("ALTER TABLE stories ADD COLUMN scene_data TEXT")
    except Exception:
        pass  # 이미 존재하면 무시

    try:
        with get_db() as conn:
            conn.execute("ALTER TABLE stories ADD COLUMN is_read INTEGER DEFAULT 0")
    except Exception:
        pass  # 이미 존재하면 무시

    # 새로운 마이그레이션: layout_data 컬럼 추가
    try:
        with get_db() as conn:
            conn.execute("ALTER TABLE stories ADD COLUMN layout_data TEXT")
    except Exception:
        pass  # 이미 존재하면 무시


# ─── CRUD Helpers ────────────────────────────────────────────────────────────


def get_random_word():
    with get_db() as conn:
        row = conn.execute("SELECT * FROM words ORDER BY RANDOM() LIMIT 1").fetchone()
        return dict(row) if row else None


def get_word(word_id):
    with get_db() as conn:
        row = conn.execute("SELECT * FROM words WHERE id=?", (word_id,)).fetchone()
        return dict(row) if row else None


def get_all_words():
    with get_db() as conn:
        rows = conn.execute("SELECT * FROM words ORDER BY category, korean").fetchall()
        return [dict(r) for r in rows]



def save_drawing(word_id, image_data, file_path=None):
    with get_db() as conn:
        cursor = conn.execute(
            "INSERT INTO drawings (word_id, image_data, file_path) VALUES (?, ?, ?)",
            (word_id, image_data, file_path),
        )
        return cursor.lastrowid


def get_all_drawings():
    with get_db() as conn:
        rows = conn.execute("""
            SELECT d.id, d.word_id, d.file_path, d.created_at,
                   w.korean, w.english, w.emoji, w.category
            FROM drawings d JOIN words w ON d.word_id = w.id
            ORDER BY w.korean ASC
        """).fetchall()
        return [dict(r) for r in rows]


def delete_drawing(drawing_id):
    with get_db() as conn:
        conn.execute("DELETE FROM drawings WHERE id=?", (drawing_id,))


def get_drawing_by_word_id(word_id):
    with get_db() as conn:
        row = conn.execute(
            "SELECT * FROM drawings WHERE word_id=? ORDER BY created_at DESC LIMIT 1",
            (word_id,),
        ).fetchone()
        return dict(row) if row else None


def get_drawing_with_data(drawing_id):
    with get_db() as conn:
        row = conn.execute(
            """
            SELECT d.*, w.korean, w.english, w.emoji
            FROM drawings d JOIN words w ON d.word_id = w.id
            WHERE d.id=?
        """,
            (drawing_id,),
        ).fetchone()
        return dict(row) if row else None


def get_drawings_by_ids(drawing_ids):
    if not drawing_ids:
        return []
    placeholders = ",".join("?" * len(drawing_ids))
    with get_db() as conn:
        rows = conn.execute(
            f"""
            SELECT d.id, d.word_id, d.image_data, d.file_path, d.created_at,
                   w.korean, w.english, w.emoji
            FROM drawings d LEFT JOIN words w ON d.word_id = w.id
            WHERE d.id IN ({placeholders})
        """,
            drawing_ids,
        ).fetchall()
        # 선택 순서(drawing_ids 순서)를 유지하여 반환
        rows_map = {row["id"]: dict(row) for row in rows}
        return [rows_map[i] for i in drawing_ids if i in rows_map]


def save_story(
    title,
    content,
    moral,
    keywords,
    drawing_ids,
    illustration_path=None,
    audio_path=None,
    scene_data=None,
    layout_data=None,
):
    with get_db() as conn:
        cursor = conn.execute(
            """INSERT INTO stories
               (title, content, moral, keywords, drawing_ids, illustration_path, audio_path, scene_data, layout_data)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                title,
                content,
                moral,
                json.dumps(keywords, ensure_ascii=False),
                json.dumps(drawing_ids),
                illustration_path,
                audio_path,
                json.dumps(scene_data, ensure_ascii=False) if scene_data else None,
                json.dumps(layout_data, ensure_ascii=False) if layout_data else None,
            ),
        )
        return cursor.lastrowid


def get_story(story_id):
    with get_db() as conn:
        row = conn.execute("SELECT * FROM stories WHERE id=?", (story_id,)).fetchone()
        if not row:
            return None
        story = dict(row)
        story["keywords"] = json.loads(story["keywords"])
        story["drawing_ids"] = json.loads(story["drawing_ids"])
        story["scene_data"] = (
            json.loads(story["scene_data"]) if story.get("scene_data") else None
        )
        story["layout_data"] = (
            json.loads(story["layout_data"]) if story.get("layout_data") else None
        )
        return story


def get_all_stories():
    with get_db() as conn:
        rows = conn.execute("SELECT * FROM stories ORDER BY created_at DESC").fetchall()
        result = []
        for row in rows:
            s = dict(row)
            s["keywords"] = json.loads(s["keywords"])
            s["drawing_ids"] = json.loads(s["drawing_ids"])
            s["scene_data"] = (
                json.loads(s["scene_data"]) if s.get("scene_data") else None
            )
            s["layout_data"] = (
                json.loads(s["layout_data"]) if s.get("layout_data") else None
            )
            result.append(s)
        return result


def delete_story(story_id):
    with get_db() as conn:
        conn.execute("DELETE FROM stories WHERE id=?", (story_id,))


def mark_story_read(story_id):
    with get_db() as conn:
        conn.execute("UPDATE stories SET is_read = 1 WHERE id=?", (story_id,))


def get_story_count():
    """전체 생성된 동화책의 개수를 반환합니다."""
    with get_db() as conn:
        row = conn.execute("SELECT COUNT(*) FROM stories").fetchone()
        return row[0] if row else 0