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

try:
    from config import mysql_host, mysql_user, mysql_password, mysql_database
except ImportError:
    mysql_host = os.environ.get('MYSQL_HOST', 'localhost')
    mysql_user = os.environ.get('MYSQL_USER', 'root')
    mysql_password = os.environ.get('MYSQL_PASSWORD', '')
    mysql_database = os.environ.get('MYSQL_DATABASE', 'maixiang_ai')


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)
                )
            """)
            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()
