from db.connection import get_connection
from loguru import logger

def create_tables():
    try:
        with get_connection() as conn:
            with conn.cursor() as cur:
                # Пользователи
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS users (
                        id SERIAL PRIMARY KEY,
                        email VARCHAR(255) UNIQUE NOT NULL,
                        password_hash VARCHAR(255) NOT NULL,
                        name VARCHAR(255) NOT NULL,
                        role VARCHAR(50) NOT NULL DEFAULT 'viewer',
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    )
                """)

                # Роботы
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS robots (
                        id VARCHAR(50) PRIMARY KEY,
                        status VARCHAR(50) DEFAULT 'active',
                        battery_level INTEGER DEFAULT 100,
                        last_update TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        current_zone VARCHAR(10),
                        current_row INTEGER,
                        current_shelf INTEGER
                    )
                """)

                # Товары
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS products (
                        id VARCHAR(50) PRIMARY KEY,
                        name VARCHAR(255) NOT NULL,
                        category VARCHAR(100),
                        min_stock INTEGER DEFAULT 10,
                        optimal_stock INTEGER DEFAULT 100
                    )
                """)

                # История инвентаризации
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS inventory_history (
                        id SERIAL PRIMARY KEY,
                        robot_id VARCHAR(50) REFERENCES robots(id),
                        product_id VARCHAR(50) REFERENCES products(id),
                        quantity INTEGER NOT NULL,
                        zone VARCHAR(10) NOT NULL,
                        row_number INTEGER,
                        shelf_number INTEGER,
                        status VARCHAR(50),
                        scanned_at TIMESTAMP NOT NULL,
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    )
                """)

                # Прогнозы ИИ
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS ai_predictions (
                        id SERIAL PRIMARY KEY,
                        product_id VARCHAR(50) REFERENCES products(id),
                        prediction_date DATE NOT NULL,
                        days_until_stockout INTEGER,
                        recommended_order INTEGER,
                        confidence_score DECIMAL(3,2),
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    )
                """)

                conn.commit()
                logger.debug("Все таблицы успешно созданы")

    except Exception as e:
        logger.error(f"Ошибка создания таблиц: {e}")
        raise


def create_indexes():
    try:
        with get_connection() as conn:
            with conn.cursor() as cur:
                cur.execute("CREATE INDEX IF NOT EXISTS idx_inventory_scanned ON inventory_history(scanned_at DESC)")
                cur.execute("CREATE INDEX IF NOT EXISTS idx_inventory_product ON inventory_history(product_id)")
                cur.execute("CREATE INDEX IF NOT EXISTS idx_inventory_zone ON inventory_history(zone)")

                conn.commit()
                logger.debug("Индексы созданы")

    except Exception as e:
        logger.error(f"Ошибка создания индексов: {e}")
        raise


def insert_sample_data():
    try:
        with get_connection() as conn:
            with conn.cursor() as cur:
                cur.execute("""
                    INSERT INTO users (email, password_hash, name, role) 
                    VALUES 
                    ('admin@warehouse.com', 'hash1', 'Администратор', 'admin'),
                    ('operator@warehouse.com', 'hash2', 'Оператор Иванов', 'operator'),
                    ('viewer@warehouse.com', 'hash3', 'Наблюдатель Петров', 'viewer')
                    ON CONFLICT (email) DO NOTHING
                """)

                cur.execute("""
                    INSERT INTO robots (id, status, battery_level, current_zone) 
                    VALUES 
                    ('RB-001', 'active', 85, 'A'),
                    ('RB-002', 'active', 45, 'B'), 
                    ('RB-003', 'maintenance', 100, NULL)
                    ON CONFLICT (id) DO NOTHING
                """)

                cur.execute("""
                    INSERT INTO products (id, name, category, min_stock, optimal_stock) 
                    VALUES 
                    ('TEL-1234', 'Смартфон X', 'Электроника', 5, 50),
                    ('NOTE-567', 'Ноутбук Pro', 'Электроника', 3, 20),
                    ('ACC-999', 'Чехол для телефона', 'Аксессуары', 10, 100)
                    ON CONFLICT (id) DO NOTHING
                """)

                conn.commit()
                logger.debug("Тестовые данные добавлены")

    except Exception as e:
        logger.error(f"Ошибка добавления тестовых данных: {e}")
        raise


def initialize_database():
    logger.info("Начинаем инициализацию базы данных...")

    create_tables()
    create_indexes()
    insert_sample_data()

    logger.debug("База данных успешно инициализирована!")


if __name__ == "__main__":
    initialize_database()