#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
커스텀 처방 관리를 위한 데이터베이스 스키마 업데이트
"""

import sqlite3
from datetime import datetime

def get_connection():
    """데이터베이스 연결"""
    return sqlite3.connect('database/kdrug.db')

def add_custom_fields():
    """커스텀 처방 관련 필드 추가"""
    conn = get_connection()
    cursor = conn.cursor()

    print("\n" + "="*60)
    print("커스텀 처방 관리를 위한 DB 스키마 업데이트")
    print("="*60)

    try:
        # 1. compounds 테이블에 커스텀 관련 필드 추가
        print("\n1. compounds 테이블 업데이트...")

        # is_custom 컬럼 추가
        cursor.execute("""
            ALTER TABLE compounds
            ADD COLUMN is_custom BOOLEAN DEFAULT 0
        """)
        print("   ✓ is_custom 컬럼 추가")

        # custom_summary 컬럼 추가
        cursor.execute("""
            ALTER TABLE compounds
            ADD COLUMN custom_summary TEXT
        """)
        print("   ✓ custom_summary 컬럼 추가")

        # custom_type 컬럼 추가
        cursor.execute("""
            ALTER TABLE compounds
            ADD COLUMN custom_type TEXT
        """)
        print("   ✓ custom_type 컬럼 추가")

    except sqlite3.OperationalError as e:
        if "duplicate column" in str(e):
            print("   ⚠ 이미 컬럼이 존재합니다.")
        else:
            raise e

    try:
        # 2. compound_ingredients 테이블에 modification 관련 필드 추가
        print("\n2. compound_ingredients 테이블 업데이트...")

        # modification_type 컬럼 추가
        cursor.execute("""
            ALTER TABLE compound_ingredients
            ADD COLUMN modification_type TEXT DEFAULT 'original'
        """)
        print("   ✓ modification_type 컬럼 추가")

        # original_grams 컬럼 추가 (원래 용량 저장)
        cursor.execute("""
            ALTER TABLE compound_ingredients
            ADD COLUMN original_grams REAL
        """)
        print("   ✓ original_grams 컬럼 추가")

    except sqlite3.OperationalError as e:
        if "duplicate column" in str(e):
            print("   ⚠ 이미 컬럼이 존재합니다.")
        else:
            raise e

    # 3. 인덱스 추가
    try:
        print("\n3. 인덱스 생성...")
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_compounds_is_custom
            ON compounds(is_custom)
        """)
        print("   ✓ is_custom 인덱스 생성")

        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_compounds_patient_custom
            ON compounds(patient_id, is_custom)
        """)
        print("   ✓ patient_id + is_custom 복합 인덱스 생성")

    except sqlite3.OperationalError as e:
        print(f"   ⚠ 인덱스 생성 중 오류: {e}")

    conn.commit()

    # 4. 스키마 확인
    print("\n4. 업데이트된 스키마 확인...")

    cursor.execute("PRAGMA table_info(compounds)")
    columns = cursor.fetchall()
    print("\n   compounds 테이블 컬럼:")
    for col in columns:
        if col[1] in ['is_custom', 'custom_summary', 'custom_type']:
            print(f"     ✓ {col[1]:20s} {col[2]}")

    cursor.execute("PRAGMA table_info(compound_ingredients)")
    columns = cursor.fetchall()
    print("\n   compound_ingredients 테이블 컬럼:")
    for col in columns:
        if col[1] in ['modification_type', 'original_grams']:
            print(f"     ✓ {col[1]:20s} {col[2]}")

    conn.close()

    print("\n" + "="*60)
    print("✅ DB 스키마 업데이트 완료!")
    print("="*60)

def test_custom_fields():
    """업데이트된 필드 테스트"""
    conn = get_connection()
    cursor = conn.cursor()

    print("\n테스트: 커스텀 필드 동작 확인...")

    try:
        # 테스트 쿼리
        cursor.execute("""
            SELECT
                compound_id,
                is_custom,
                custom_summary,
                custom_type
            FROM compounds
            LIMIT 1
        """)

        result = cursor.fetchone()
        if result:
            print("   ✓ compounds 테이블 커스텀 필드 정상")
        else:
            print("   ℹ compounds 테이블이 비어있습니다.")

        cursor.execute("""
            SELECT
                compound_ingredient_id,
                modification_type,
                original_grams
            FROM compound_ingredients
            LIMIT 1
        """)

        result = cursor.fetchone()
        if result:
            print("   ✓ compound_ingredients 테이블 커스텀 필드 정상")
        else:
            print("   ℹ compound_ingredients 테이블이 비어있습니다.")

    except Exception as e:
        print(f"   ✗ 테스트 실패: {e}")

    conn.close()

def main():
    """메인 실행"""
    add_custom_fields()
    test_custom_fields()

if __name__ == "__main__":
    main()