"""
多店铺销售预测系统 - 数据库初始化脚本
创建店铺相关表结构
"""

import sqlite3
import os
from datetime import datetime

def init_multi_store_database(db_path='prediction_history.db'):
    """初始化多店铺数据库结构"""
    
    try:
        # 连接数据库
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        print("开始初始化多店铺数据库结构...")
        
        # 1. 创建店铺表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS stores (
                store_id VARCHAR(20) PRIMARY KEY,
                store_name VARCHAR(100) NOT NULL,
                location VARCHAR(200),
                size FLOAT,
                type VARCHAR(50),
                opening_date DATE,
                status VARCHAR(20) DEFAULT 'active',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        print("店铺表(stores)创建成功")
        
        # 2. 创建产品表（如果不存在）
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS products (
                product_id VARCHAR(20) PRIMARY KEY,
                product_name VARCHAR(100) NOT NULL,
                category VARCHAR(50),
                price FLOAT,
                unit VARCHAR(20),
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        print("产品表(products)创建成功")
        
        # 3. 创建店铺-产品关联表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS store_products (
                store_id VARCHAR(20),
                product_id VARCHAR(20),
                first_sale_date DATE,
                is_active BOOLEAN DEFAULT 1,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (store_id, product_id),
                FOREIGN KEY (store_id) REFERENCES stores(store_id),
                FOREIGN KEY (product_id) REFERENCES products(product_id)
            )
        ''')
        print("店铺-产品关联表(store_products)创建成功")
        
        # 4. 检查是否需要修改现有表
        cursor.execute("PRAGMA table_info(prediction_history)")
        columns = [column[1] for column in cursor.fetchall()]
        
        if 'store_id' not in columns:
            # 为预测历史表添加store_id字段
            cursor.execute('ALTER TABLE prediction_history ADD COLUMN store_id VARCHAR(20)')
            cursor.execute('ALTER TABLE prediction_history ADD COLUMN store_name VARCHAR(100)')
            print("预测历史表已添加店铺字段")
        else:
            print("预测历史表已包含店铺字段")
        
        # 5. 插入示例店铺数据
        sample_stores = [
            ('S001', '旗舰店-市中心', '市中心商业区', 200.0, 'flagship', '2020-01-01', 'active'),
            ('S002', '标准店-东区', '东区购物中心', 150.0, 'standard', '2020-03-15', 'active'),
            ('S003', '社区店-南区', '南区居民区', 80.0, 'community', '2020-06-01', 'active'),
            ('S004', '标准店-西区', '西区商业街', 120.0, 'standard', '2020-09-10', 'active'),
            ('S005', '社区店-北区', '北区社区中心', 90.0, 'community', '2021-01-20', 'active')
        ]
        
        for store in sample_stores:
            cursor.execute('''
                INSERT OR IGNORE INTO stores 
                (store_id, store_name, location, size, type, opening_date, status) 
                VALUES (?, ?, ?, ?, ?, ?, ?)
            ''', store)
        print("示例店铺数据插入成功")
        
        # 6. 插入示例产品数据
        sample_products = [
            ('P001', '感冒灵颗粒', '感冒药', 15.80, '盒'),
            ('P002', '布洛芬片', '解热镇痛', 12.50, '盒'),
            ('P003', '阿莫西林胶囊', '抗生素', 18.90, '盒'),
            ('P004', '维生素C片', '维生素', 8.60, '瓶'),
            ('P005', '板蓝根颗粒', '清热解毒', 13.20, '盒')
        ]
        
        for product in sample_products:
            cursor.execute('''
                INSERT OR IGNORE INTO products 
                (product_id, product_name, category, price, unit) 
                VALUES (?, ?, ?, ?, ?)
            ''', product)
        print("示例产品数据插入成功")
        
        # 7. 创建店铺-产品关联
        # 为每个店铺关联所有产品
        for store in sample_stores:
            store_id = store[0]
            for product in sample_products:
                product_id = product[0]
                cursor.execute('''
                    INSERT OR IGNORE INTO store_products 
                    (store_id, product_id, first_sale_date, is_active) 
                    VALUES (?, ?, ?, ?)
                ''', (store_id, product_id, '2020-01-01', 1))
        print("店铺-产品关联数据创建成功")
        
        # 8. 创建索引以提高查询性能
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_stores_status ON stores(status)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_stores_type ON stores(type)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_store_products_store ON store_products(store_id)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_store_products_product ON store_products(product_id)')
        print("数据库索引创建成功")
        
        # 提交更改
        conn.commit()
        conn.close()
        
        print("\n多店铺数据库初始化完成！")
        print("店铺表结构:")
        print("- stores: 店铺基本信息")
        print("- products: 产品信息")
        print("- store_products: 店铺-产品关联")
        print("- prediction_history: 预测历史（已添加店铺字段）")
        
        return True
        
    except Exception as e:
        print(f"数据库初始化失败: {e}")
        return False

def get_db_connection(db_path='prediction_history.db'):
    """获取数据库连接"""
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row  # 让查询结果可以通过列名访问
    return conn

def check_database_structure():
    """检查数据库结构"""
    try:
        conn = get_db_connection()
        cursor = conn.cursor()
        
        print("当前数据库表结构:")
        
        # 获取所有表
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
        tables = cursor.fetchall()
        
        for table in tables:
            table_name = table['name']
            print(f"\n表: {table_name}")
            
            cursor.execute(f"PRAGMA table_info({table_name})")
            columns = cursor.fetchall()
            
            for column in columns:
                print(f"  - {column['name']}: {column['type']}")
        
        conn.close()
        
    except Exception as e:
        print(f"检查数据库结构失败: {e}")

if __name__ == "__main__":
    # 切换到server目录
    script_dir = os.path.dirname(os.path.abspath(__file__))
    os.chdir(script_dir)
    
    print("多店铺销售预测系统 - 数据库初始化")
    print("=" * 50)
    
    # 初始化数据库
    success = init_multi_store_database()
    
    if success:
        print("\n" + "=" * 50)
        check_database_structure()
    else:
        print("数据库初始化失败，请检查错误信息")