"""
数据库连接与初始化
"""
import os
import pymysql
from werkzeug.security import generate_password_hash

try:
    from config import mysql_host as config_mysql_host
    from config import mysql_user as config_mysql_user
    from config import mysql_password as config_mysql_password
    from config import mysql_database as config_mysql_database
except ImportError:
    config_mysql_host = 'localhost'
    config_mysql_user = 'root'
    config_mysql_password = ''
    config_mysql_database = 'maixiang_ai'

mysql_host = os.environ.get('MYSQL_HOST', config_mysql_host)
mysql_user = os.environ.get('MYSQL_USER', config_mysql_user)
mysql_password = os.environ.get('MYSQL_PASSWORD', config_mysql_password)
mysql_database = os.environ.get('MYSQL_DATABASE', config_mysql_database)


def get_db():
    return pymysql.connect(
        host=mysql_host,
        user=mysql_user,
        password=mysql_password,
        database=mysql_database,
        charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor
    )


def _create_initial_admin():
    """若没有任何管理员，则创建默认超级管理员（首次启动时，仅一个）"""
    try:
        conn = get_db()
        with conn.cursor() as cur:
            cur.execute("SELECT id FROM users WHERE role = 'super_admin' LIMIT 1")
            if cur.fetchone():
                conn.close()
                return
            admin_user = os.environ.get('ADMIN_USER', 'admin')
            admin_pwd = os.environ.get('ADMIN_PASSWORD', 'admin123')
            pwd_hash = generate_password_hash(admin_pwd, method='pbkdf2:sha256')
            cur.execute(
                "INSERT INTO users (username, password_hash, is_admin, role) VALUES (%s, %s, 1, 'super_admin')",
                (admin_user, pwd_hash)
            )
        conn.commit()
        conn.close()
    except Exception:
        pass


def init_db():
    """初始化数据库表，若不存在则创建"""
    conn = pymysql.connect(
        host=mysql_host,
        user=mysql_user,
        password=mysql_password,
        charset='utf8mb4'
    )
    try:
        with conn.cursor() as cur:
            cur.execute(f"CREATE DATABASE IF NOT EXISTS `{mysql_database}` DEFAULT CHARSET utf8mb4")
            cur.execute(f"USE `{mysql_database}`")
            cur.execute("""
                CREATE TABLE IF NOT EXISTS users (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    username VARCHAR(64) NOT NULL UNIQUE,
                    password_hash VARCHAR(256) NOT NULL,
                    is_admin TINYINT(1) DEFAULT 0,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            """)
            cur.execute("""
                CREATE TABLE IF NOT EXISTS image_history (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    user_id INT NOT NULL,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    panel_type VARCHAR(64) DEFAULT '',
                    original_urls JSON,
                    params JSON,
                    result_urls JSON,
                    long_image_url VARCHAR(1024) DEFAULT NULL,
                    INDEX idx_user_created (user_id, created_at DESC)
                )
            """)
            try:
                cur.execute("ALTER TABLE image_history ADD COLUMN long_image_url VARCHAR(1024) DEFAULT NULL")
            except Exception:
                pass
            try:
                cur.execute("ALTER TABLE users ADD COLUMN machine VARCHAR(64) DEFAULT NULL")
            except Exception:
                pass
            try:
                cur.execute("ALTER TABLE users ADD COLUMN role VARCHAR(20) DEFAULT 'normal'")
            except Exception:
                pass
            try:
                cur.execute("ALTER TABLE users ADD COLUMN created_by_id INT NULL")
            except Exception:
                pass
            cur.execute("""
                CREATE TABLE IF NOT EXISTS web_config (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    version VARCHAR(64) NOT NULL,
                    file_url VARCHAR(1024) NOT NULL,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            """)
            cur.execute("""
                CREATE TABLE IF NOT EXISTS columns (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    name VARCHAR(128) NOT NULL COMMENT '栏目名',
                    column_key VARCHAR(64) NOT NULL COMMENT '栏目标识',
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE KEY uk_column_key (column_key)
                )
            """)
            try:
                cur.execute("ALTER TABLE columns ADD COLUMN menu_type VARCHAR(20) NOT NULL DEFAULT 'app' COMMENT '菜单类型: app/admin' AFTER column_key")
            except Exception:
                pass
            try:
                cur.execute("ALTER TABLE columns ADD COLUMN route_path VARCHAR(255) NOT NULL DEFAULT '' COMMENT '菜单路由或页面标识' AFTER menu_type")
            except Exception:
                pass
            try:
                cur.execute("ALTER TABLE columns ADD COLUMN sort_order INT NOT NULL DEFAULT 0 COMMENT '菜单排序' AFTER route_path")
            except Exception:
                pass
            try:
                cur.execute("UPDATE columns SET menu_type = 'app' WHERE menu_type IS NULL OR menu_type = ''")
            except Exception:
                pass
            try:
                cur.execute("UPDATE columns SET sort_order = id WHERE sort_order IS NULL OR sort_order = 0")
            except Exception:
                pass
            cur.execute("""
                CREATE TABLE IF NOT EXISTS user_column_permission (
                    user_id INT NOT NULL,
                    column_id INT NOT NULL,
                    PRIMARY KEY (user_id, column_id),
                    FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE,
                    FOREIGN KEY (column_id) REFERENCES columns(id) ON DELETE CASCADE
                )
            """)
            try:
                cur.execute("UPDATE users SET role = 'normal' WHERE (role IS NULL OR role = '') AND (is_admin = 0 OR is_admin IS NULL)")
                cur.execute("SELECT MIN(id) AS mid FROM users WHERE is_admin = 1")
                row = cur.fetchone()
                if row and row.get('mid'):
                    mid = row['mid']
                    cur.execute("UPDATE users SET role = 'super_admin' WHERE id = %s", (mid,))
                    cur.execute("UPDATE users SET role = 'admin' WHERE is_admin = 1 AND id != %s", (mid,))
                    cur.execute("UPDATE users SET created_by_id = %s WHERE role = 'admin' AND (created_by_id IS NULL)", (mid,))
            except Exception:
                pass
        conn.commit()
    finally:
        conn.close()
    _create_initial_admin()