import psycopg2
import os
import base64
import hashlib
import hmac
import secrets
from psycopg2.extras import RealDictCursor

# ---------------- 数据库配置区 ----------------
DB_CONFIG = {
    "dbname": os.getenv("POSTGRES_DB", "cnc_eval_db"),
    "user": os.getenv("POSTGRES_USER", "main"),
    "password": os.getenv("POSTGRES_PASSWORD", ""),
    "host": os.getenv("POSTGRES_HOST", "localhost"),
    "port": os.getenv("POSTGRES_PORT", "5433"),
}


PASSWORD_ITERATIONS = 600_000


def hash_password(password: str) -> str:
    salt = secrets.token_bytes(16)
    digest = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt, PASSWORD_ITERATIONS)
    return "pbkdf2_sha256${}${}${}".format(
        PASSWORD_ITERATIONS,
        base64.b64encode(salt).decode("ascii"),
        base64.b64encode(digest).decode("ascii"),
    )


def verify_password(password: str, stored_hash: str) -> bool:
    if stored_hash.startswith("pbkdf2_sha256$"):
        try:
            _, iterations, salt, expected = stored_hash.split("$", 3)
            digest = hashlib.pbkdf2_hmac(
                "sha256",
                password.encode("utf-8"),
                base64.b64decode(salt),
                int(iterations),
            )
            return hmac.compare_digest(base64.b64encode(digest).decode("ascii"), expected)
        except (ValueError, TypeError):
            return False
    legacy = hashlib.sha256(password.encode("utf-8")).hexdigest()
    return hmac.compare_digest(legacy, stored_hash)


def password_needs_rehash(stored_hash: str) -> bool:
    return not stored_hash.startswith(f"pbkdf2_sha256${PASSWORD_ITERATIONS}$")

def get_db_connection():
    """获取数据库连接的通用函数"""
    return psycopg2.connect(**DB_CONFIG)


def ensure_user_table():
    conn = get_db_connection()
    cursor = conn.cursor()
    try:
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS sys_users (
                id SERIAL PRIMARY KEY,
                username VARCHAR(50) UNIQUE NOT NULL,
                password_hash VARCHAR(200) NOT NULL,
                real_name VARCHAR(50),
                role VARCHAR(20) DEFAULT 'user',
                email VARCHAR(120),
                is_active BOOLEAN DEFAULT TRUE,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                last_login TIMESTAMP
            )
        """)
        cursor.execute("ALTER TABLE sys_users ADD COLUMN IF NOT EXISTS real_name VARCHAR(50)")
        cursor.execute("ALTER TABLE sys_users ADD COLUMN IF NOT EXISTS role VARCHAR(20) DEFAULT 'user'")
        cursor.execute("ALTER TABLE sys_users ADD COLUMN IF NOT EXISTS email VARCHAR(120)")
        cursor.execute("ALTER TABLE sys_users ADD COLUMN IF NOT EXISTS is_active BOOLEAN DEFAULT TRUE")
        cursor.execute("ALTER TABLE sys_users ADD COLUMN IF NOT EXISTS created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP")
        cursor.execute("ALTER TABLE sys_users ADD COLUMN IF NOT EXISTS last_login TIMESTAMP")
        cursor.execute("SELECT 1 FROM sys_users WHERE username = 'admin'")
        admin_exists = cursor.fetchone() is not None
        initial_password = os.getenv("INITIAL_ADMIN_PASSWORD")
        if not admin_exists and not initial_password:
            raise RuntimeError("首次启动必须配置 INITIAL_ADMIN_PASSWORD")
        if admin_exists:
            conn.commit()
            return
        cursor.execute("""
            INSERT INTO sys_users (username, password_hash, real_name, role, email, is_active)
            SELECT %s, %s, %s, %s, %s, %s
            WHERE NOT EXISTS (SELECT 1 FROM sys_users WHERE username = %s)
        """, (
            'admin',
            hash_password(initial_password),
            '系统管理员',
            'admin',
            'admin@precision.local',
            True,
            'admin'
        ))
        conn.commit()
    finally:
        cursor.close()
        conn.close()


def ensure_metric_bind_columns():
    conn = get_db_connection()
    cursor = conn.cursor()
    try:
        cursor.execute("ALTER TABLE tb_asset_metric_bind ADD COLUMN IF NOT EXISTS lower_limit FLOAT")
        cursor.execute("ALTER TABLE tb_asset_metric_bind ADD COLUMN IF NOT EXISTS baseline_value FLOAT")
        cursor.execute("ALTER TABLE tb_asset_metric_bind ADD COLUMN IF NOT EXISTS failure_value FLOAT")
        cursor.execute("ALTER TABLE tb_asset_metric_bind ADD COLUMN IF NOT EXISTS ideal_value FLOAT")
        cursor.execute("ALTER TABLE tb_asset_metric_bind ADD COLUMN IF NOT EXISTS extreme_value FLOAT")
        cursor.execute("ALTER TABLE tb_asset_metric_bind ADD COLUMN IF NOT EXISTS criterion_weight FLOAT")
        cursor.execute("ALTER TABLE tb_asset_metric_bind ADD COLUMN IF NOT EXISTS indicator_weight FLOAT")
        cursor.execute("ALTER TABLE tb_asset_metric_bind ADD COLUMN IF NOT EXISTS cycle_hours FLOAT")
        conn.commit()
    finally:
        cursor.close()
        conn.close()


def ensure_product_tables():
    conn = get_db_connection()
    cursor = conn.cursor()
    try:
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS tb_product_def (
                product_id SERIAL PRIMARY KEY,
                product_code VARCHAR(50) UNIQUE NOT NULL,
                product_name VARCHAR(120) NOT NULL,
                product_detail TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS tb_asset_product_bind (
                bind_id SERIAL PRIMARY KEY,
                asset_id INTEGER NOT NULL REFERENCES tb_asset(asset_id) ON DELETE CASCADE,
                product_id INTEGER NOT NULL REFERENCES tb_product_def(product_id) ON DELETE CASCADE,
                total_quantity INTEGER,
                process_note TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(asset_id, product_id)
            )
        """)
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS tb_product_metric_config (
                config_id SERIAL PRIMARY KEY,
                product_id INTEGER NOT NULL REFERENCES tb_product_def(product_id) ON DELETE CASCADE,
                metric_id INTEGER NOT NULL REFERENCES tb_metric_dict(metric_id) ON DELETE CASCADE,
                upper_limit FLOAT,
                lower_limit FLOAT,
                ideal_value FLOAT,
                extreme_value FLOAT,
                unit VARCHAR(20),
                position_note TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(product_id, metric_id)
            )
        """)
        cursor.execute("ALTER TABLE tb_product_metric_config ADD COLUMN IF NOT EXISTS position_note TEXT")
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS tb_product_record (
                product_record_id SERIAL PRIMARY KEY,
                asset_id INTEGER NOT NULL REFERENCES tb_asset(asset_id) ON DELETE CASCADE,
                product_id INTEGER NOT NULL REFERENCES tb_product_def(product_id) ON DELETE CASCADE,
                batch_no VARCHAR(120) NOT NULL,
                sample_count INTEGER,
                production_quantity INTEGER,
                created_at TIMESTAMP NOT NULL,
                remark TEXT
            )
        """)
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS tb_product_sample_value (
                sample_id SERIAL PRIMARY KEY,
                record_id INTEGER NOT NULL REFERENCES tb_product_record(product_record_id) ON DELETE CASCADE,
                metric_id INTEGER NOT NULL REFERENCES tb_metric_dict(metric_id) ON DELETE CASCADE,
                sample_no VARCHAR(40) NOT NULL,
                value FLOAT NOT NULL
            )
        """)
        conn.commit()
    finally:
        cursor.close()
        conn.close()


def ensure_interchangeability_metrics():
    conn = get_db_connection()
    cursor = conn.cursor()
    try:
        cursor.execute(
            """
            UPDATE tb_metric_dict
            SET category_level_1 = '加工精度'
            WHERE category_level_1 = '互换性指标'
            """
        )
        metric_category_map = {
            "线性尺寸偏差": "尺寸",
            "角度尺寸偏差": "尺寸",
            "直线度": "直线度",
            "平面度": "平面度",
            "圆度": "圆度",
            "圆柱度": "圆柱度",
            "线轮廓度": "线轮廓度",
            "面轮廓度": "面轮廓度",
            "平行度": "平行度",
            "垂直度": "垂直度",
            "倾斜度": "倾斜度",
            "位置度": "位置度",
            "同轴度": "同轴度",
            "同心度": "同心度",
            "对称度": "对称度",
            "圆跳动": "圆跳动",
            "全跳动": "全跳动",
            "表面粗糙度": "粗糙度",
            "表面波纹度": "波纹度",
        }
        for metric_name, category_level_2 in metric_category_map.items():
            cursor.execute(
                """
                UPDATE tb_metric_dict
                SET category_level_1 = '加工精度', category_level_2 = %s
                WHERE category_level_1 = '加工精度'
                  AND metric_name = %s
                  AND category_level_2 IN ('尺寸精度', '形状公差', '方向公差', '位置公差', '跳动公差', '表面结构')
                """,
                (category_level_2, metric_name),
            )
        metrics = [
            ("尺寸", "线性尺寸偏差", "抽样数组", "加工精度制品指标；可在产品配置中按具体作用部位补充说明。"),
            ("尺寸", "角度尺寸偏差", "抽样数组", "加工精度制品指标；可在产品配置中按具体作用部位补充说明。"),
            ("直线度", "直线度", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("平面度", "平面度", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("圆度", "圆度", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("圆柱度", "圆柱度", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("线轮廓度", "线轮廓度", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("面轮廓度", "面轮廓度", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("平行度", "平行度", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("垂直度", "垂直度", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("倾斜度", "倾斜度", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("位置度", "位置度", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("同轴度", "同轴度", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("同心度", "同心度", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("对称度", "对称度", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("圆跳动", "圆跳动", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("全跳动", "全跳动", "抽样数组", "加工精度制品指标；GB/T 1182-2018 几何公差项目。"),
            ("粗糙度", "表面粗糙度", "抽样数组", "加工精度制品指标；表面结构/粗糙度评价参数。"),
            ("波纹度", "表面波纹度", "抽样数组", "加工精度制品指标；表面结构评价参数。"),
        ]
        for category_level_2, metric_name, value_type, description in metrics:
            cursor.execute(
                """
                INSERT INTO tb_metric_dict (category_level_1, category_level_2, metric_name, value_type, description)
                SELECT %s, %s, %s, %s, %s
                WHERE NOT EXISTS (
                    SELECT 1
                    FROM tb_metric_dict
                    WHERE category_level_1 = %s AND category_level_2 = %s AND metric_name = %s
                )
                """,
                (
                    "加工精度",
                    category_level_2,
                    metric_name,
                    value_type,
                    description,
                    "加工精度",
                    category_level_2,
                    metric_name,
                ),
            )
        conn.commit()
    finally:
        cursor.close()
        conn.close()
