"""
数据库连接和用户认证
"""
import psycopg2
from psycopg2 import pool
from psycopg2.extras import RealDictCursor
from werkzeug.security import generate_password_hash, check_password_hash
from flask import session
import functools
from config import Config

_pool = None


def init_db_pool():
    """初始化数据库连接池（应用启动时调用一次）"""
    global _pool
    if _pool is not None:
        return
    try:
        _pool = pool.ThreadedConnectionPool(
            minconn=2,
            maxconn=20,
            host=Config.DB_HOST,
            port=Config.DB_PORT,
            database=Config.DB_NAME,
            user=Config.DB_USER,
            password=Config.DB_PASSWORD, sslmode='disable',
        )
        print(f"[DB] 连接池初始化成功 (min=2, max=20)")
    except Exception as e:
        print(f"[DB] 连接池初始化失败，降级为逐次建连: {e}")
        _pool = None


def get_db():
    """从连接池获取数据库连接（用完必须调用 put_db(conn) 归还）"""
    global _pool
    if _pool is not None:
        try:
            conn = _pool.getconn()
            if conn and not conn.closed:
                conn.autocommit = False
                return conn
        except Exception as e:
            print(f"[DB] 连接池获取失败: {e}")
    try:
        conn = psycopg2.connect(
            host=Config.DB_HOST,
            port=Config.DB_PORT,
            database=Config.DB_NAME,
            user=Config.DB_USER,
            password=Config.DB_PASSWORD, sslmode="disable"
        )
        return conn
    except Exception as e:
        print(f"数据库连接失败: {e}")
        return None


def put_db(conn):
    """归还连接到连接池"""
    if conn is None:
        return
    global _pool
    if _pool is not None:
        try:
            _pool.putconn(conn)
            return
        except Exception:
            pass
    try:
        conn.close()
    except Exception:
        pass


def login_required(f):
    """登录验证装饰器"""
    @functools.wraps(f)
    def decorated_function(*args, **kwargs):
        if 'user_id' not in session:
            from flask import jsonify
            return jsonify({'success': False, 'error': '请先登录'}), 401
        return f(*args, **kwargs)
    return decorated_function


def get_current_user_id():
    """获取当前登录用户ID"""
    return session.get('user_id')


def get_current_username():
    """获取当前登录用户名"""
    return session.get('username')


# ========== 可用资金操作 ==========

def db_get_available_cash(user_id):
    """获取用户可用资金"""
    conn = get_db()
    if not conn:
        return 0
    
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        cur.execute("SELECT available_cash FROM users WHERE id = %s", (user_id,))
        result = cur.fetchone()
        return float(result['available_cash'] or 0) if result else 0
    finally:
        put_db(conn)


def db_update_available_cash(user_id, amount):
    """更新用户可用资金"""
    conn = get_db()
    if not conn:
        return False, '数据库连接失败'
    
    try:
        cur = conn.cursor()
        cur.execute("UPDATE users SET available_cash = %s WHERE id = %s", (amount, user_id))
        conn.commit()
        return True, None
    except Exception as e:
        conn.rollback()
        return False, str(e)
    finally:
        put_db(conn)


# ========== 用户操作 ==========

def create_user(email, password):
    """创建用户（使用邮箱）"""
    conn = get_db()
    if not conn:
        return None, '数据库连接失败'
    
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        
        # 检查邮箱是否已存在
        cur.execute("SELECT id FROM users WHERE email = %s OR username = %s", (email, email))
        if cur.fetchone():
            return None, '该邮箱已注册'
        
        # 创建用户（username和email都存邮箱）
        password_hash = generate_password_hash(password)
        cur.execute(
            "INSERT INTO users (username, email, password_hash) VALUES (%s, %s, %s) RETURNING id, username, email",
            (email, email, password_hash)
        )
        user = cur.fetchone()
        conn.commit()
        
        return user, None
    except Exception as e:
        conn.rollback()
        return None, str(e)
    finally:
        put_db(conn)


def verify_user(email, password):
    """验证用户登录（使用邮箱）"""
    conn = get_db()
    if not conn:
        return None, '数据库连接失败'
    
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        # 同时检查email和username字段（兼容旧数据）
        cur.execute("SELECT * FROM users WHERE email = %s OR username = %s", (email, email))
        user = cur.fetchone()
        
        if not user or not check_password_hash(user['password_hash'], password):
            return None, '邮箱或密码错误'
        
        return {'id': user['id'], 'username': user.get('email') or user['username']}, None
    finally:
        put_db(conn)


def change_user_password(user_id, old_password, new_password):
    """修改用户密码"""
    conn = get_db()
    if not conn:
        return False, '数据库连接失败'
    
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        cur.execute("SELECT * FROM users WHERE id = %s", (user_id,))
        user = cur.fetchone()
        
        if not user:
            return False, '用户不存在'
        
        if not check_password_hash(user['password_hash'], old_password):
            return False, '当前密码错误'
        
        new_hash = generate_password_hash(new_password)
        cur.execute("UPDATE users SET password_hash = %s WHERE id = %s", (new_hash, user_id))
        conn.commit()
        return True, None
    except Exception as e:
        return False, str(e)
    finally:
        put_db(conn)


# ========== 交易记录操作（数据库版） ==========

def db_get_trades(user_id):
    """从数据库获取用户交易记录"""
    conn = get_db()
    if not conn:
        return []
    
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        cur.execute("""
            SELECT id, stock_code, stock_name, trade_type, price, quantity,
                   trade_date::text, reason, result, profit_amount, stop_loss_price, notes,
                   created_at::text
            FROM trades 
            WHERE user_id = %s 
            ORDER BY trade_date DESC, created_at DESC
        """, (user_id,))
        return cur.fetchall()
    finally:
        put_db(conn)


def db_get_trade(user_id, trade_id):
    """获取单条交易记录"""
    conn = get_db()
    if not conn:
        return None
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        cur.execute("""
            SELECT id, stock_code, stock_name, trade_type, price, quantity,
                   trade_date::text, reason, result, profit_amount, stop_loss_price, notes,
                   created_at::text
            FROM trades WHERE id = %s AND user_id = %s
        """, (trade_id, user_id))
        return cur.fetchone()
    finally:
        put_db(conn)


def db_add_trade(user_id, data):
    """添加交易记录到数据库"""
    conn = get_db()
    if not conn:
        return None, '数据库连接失败'
    
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        cur.execute("""
            INSERT INTO trades (user_id, stock_code, stock_name, trade_type, price, 
                               quantity, trade_date, reason, result, profit_amount, 
                               stop_loss_price, notes)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            RETURNING id, stock_code, stock_name, trade_type, price, quantity,
                      trade_date::text, reason, result, profit_amount, stop_loss_price, 
                      notes, created_at::text
        """, (
            user_id,
            data.get('stock_code'),
            data.get('stock_name'),
            data.get('trade_type'),
            data.get('price'),
            data.get('quantity'),
            data.get('trade_date'),
            data.get('reason'),
            data.get('result'),
            data.get('profit_amount'),
            data.get('stop_loss_price'),
            data.get('notes')
        ))
        trade = cur.fetchone()
        conn.commit()
        return trade, None
    except Exception as e:
        conn.rollback()
        return None, str(e)
    finally:
        put_db(conn)


def db_update_trade(user_id, trade_id, data):
    """更新交易记录"""
    conn = get_db()
    if not conn:
        return None, '数据库连接失败'
    
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        cur.execute("""
            UPDATE trades SET
                stock_code = COALESCE(%s, stock_code),
                stock_name = COALESCE(%s, stock_name),
                trade_type = COALESCE(%s, trade_type),
                price = COALESCE(%s, price),
                quantity = COALESCE(%s, quantity),
                trade_date = COALESCE(%s, trade_date),
                reason = COALESCE(%s, reason),
                result = COALESCE(%s, result),
                profit_amount = COALESCE(%s, profit_amount),
                stop_loss_price = COALESCE(%s, stop_loss_price),
                notes = COALESCE(%s, notes)
            WHERE id = %s AND user_id = %s
            RETURNING id, stock_code, stock_name, trade_type, price, quantity,
                      trade_date::text, reason, result, profit_amount, stop_loss_price, 
                      notes, created_at::text
        """, (
            data.get('stock_code'),
            data.get('stock_name'),
            data.get('trade_type'),
            data.get('price'),
            data.get('quantity'),
            data.get('trade_date'),
            data.get('reason'),
            data.get('result'),
            data.get('profit_amount'),
            data.get('stop_loss_price'),
            data.get('notes'),
            trade_id,
            user_id
        ))
        trade = cur.fetchone()
        conn.commit()
        return trade, None
    except Exception as e:
        conn.rollback()
        return None, str(e)
    finally:
        put_db(conn)


def db_delete_trade(user_id, trade_id):
    """删除交易记录"""
    conn = get_db()
    if not conn:
        return False
    
    try:
        cur = conn.cursor()
        cur.execute("DELETE FROM trades WHERE id = %s AND user_id = %s", (trade_id, user_id))
        conn.commit()
        return cur.rowcount > 0
    finally:
        put_db(conn)


# ========== 关注列表操作（数据库版） ==========

def db_get_watchlist(user_id):
    """从数据库获取用户关注列表"""
    conn = get_db()
    if not conn:
        return []
    
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        cur.execute("""
            SELECT stock_code as code, stock_name as name, added_time::text
            FROM watchlist 
            WHERE user_id = %s 
            ORDER BY added_time DESC
        """, (user_id,))
        return cur.fetchall()
    finally:
        put_db(conn)


def db_add_to_watchlist(user_id, code, name):
    """添加到关注列表"""
    conn = get_db()
    if not conn:
        return None, '数据库连接失败'
    
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        cur.execute("""
            INSERT INTO watchlist (user_id, stock_code, stock_name)
            VALUES (%s, %s, %s)
            ON CONFLICT (user_id, stock_code) DO NOTHING
            RETURNING stock_code as code, stock_name as name
        """, (user_id, code, name))
        conn.commit()
        return db_get_watchlist(user_id), None
    except Exception as e:
        conn.rollback()
        return None, str(e)
    finally:
        put_db(conn)


def db_remove_from_watchlist(user_id, code):
    """从关注列表移除"""
    conn = get_db()
    if not conn:
        return None
    
    try:
        cur = conn.cursor()
        cur.execute("DELETE FROM watchlist WHERE user_id = %s AND stock_code = %s", (user_id, code))
        conn.commit()
        return db_get_watchlist(user_id)
    finally:
        put_db(conn)


# ========== 分析缓存操作（数据库版） ==========

def db_get_alerts_cache(user_id):
    """从数据库获取分析缓存"""
    conn = get_db()
    if not conn:
        return None
    
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        cur.execute("""
            SELECT data, updated_at::text as lastUpdate
            FROM alerts_cache 
            WHERE user_id = %s
        """, (user_id,))
        result = cur.fetchone()
        if result:
            return {
                'alerts': result['data'] or [],
                'lastUpdate': result['lastupdate']
            }
        return {'alerts': [], 'lastUpdate': None}
    finally:
        put_db(conn)


def db_save_alerts_cache(user_id, alerts):
    """保存分析缓存到数据库"""
    conn = get_db()
    if not conn:
        return False
    
    try:
        import json
        cur = conn.cursor()
        cur.execute("""
            INSERT INTO alerts_cache (user_id, data, updated_at)
            VALUES (%s, %s, NOW())
            ON CONFLICT (user_id) DO UPDATE SET
                data = EXCLUDED.data,
                updated_at = NOW()
        """, (user_id, json.dumps(alerts)))
        conn.commit()
        return True
    except Exception as e:
        conn.rollback()
        print(f"保存分析缓存失败: {e}")
        return False
    finally:
        put_db(conn)


# ========== 基本面数据操作（数据库版） ==========

def db_get_fundamental(code):
    """从数据库获取基本面数据（当日缓存）"""
    conn = get_db()
    if not conn:
        return None
    
    try:
        from datetime import date
        cur = conn.cursor(cursor_factory=RealDictCursor)
        cur.execute("""
            SELECT code, name, pe, pb, total_market_cap, industry, 
                   latest_price, change_pct, update_date::text, updated_at::text,
                   roe, eps, bps, revenue_yoy, profit_yoy, gross_margin, net_margin
            FROM stock_fundamental 
            WHERE code = %s AND update_date = %s
        """, (code, date.today()))
        return cur.fetchone()
    finally:
        put_db(conn)


def db_save_fundamental(code, data):
    """保存基本面数据到数据库"""
    conn = get_db()
    if not conn:
        return False
    
    try:
        from datetime import date
        cur = conn.cursor()
        cur.execute("""
            INSERT INTO stock_fundamental 
            (code, name, pe, pb, total_market_cap, industry, latest_price, change_pct, update_date,
             roe, eps, bps, revenue_yoy, profit_yoy, gross_margin, net_margin)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            ON CONFLICT (code) DO UPDATE SET
                name = EXCLUDED.name,
                pe = EXCLUDED.pe,
                pb = EXCLUDED.pb,
                total_market_cap = EXCLUDED.total_market_cap,
                industry = EXCLUDED.industry,
                latest_price = EXCLUDED.latest_price,
                change_pct = EXCLUDED.change_pct,
                update_date = EXCLUDED.update_date,
                roe = EXCLUDED.roe,
                eps = EXCLUDED.eps,
                bps = EXCLUDED.bps,
                revenue_yoy = EXCLUDED.revenue_yoy,
                profit_yoy = EXCLUDED.profit_yoy,
                gross_margin = EXCLUDED.gross_margin,
                net_margin = EXCLUDED.net_margin,
                updated_at = NOW()
        """, (
            code,
            data.get('name') or data.get('stock_name'),
            data.get('pe') or data.get('pe_ttm'),
            data.get('pb'),
            data.get('total_market_cap'),
            data.get('industry'),
            data.get('latest_price'),
            data.get('change_pct'),
            date.today(),
            data.get('roe'),
            data.get('eps'),
            data.get('bps'),
            data.get('revenue_yoy'),
            data.get('profit_yoy'),
            data.get('gross_margin'),
            data.get('net_margin'),
        ))
        conn.commit()
        return True
    except Exception as e:
        conn.rollback()
        print(f"保存基本面数据失败: {e}")
        return False
    finally:
        put_db(conn)


# ========== 资金流向历史数据操作（数据库版） ==========

def db_get_fund_flow_history(code, limit=60):
    """获取股票的资金流向历史数据"""
    conn = get_db()
    if not conn:
        return None, None
    
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        cur.execute("""
            SELECT code, trade_date::text, close_price, change_pct,
                   main_net_inflow, main_net_inflow_pct,
                   super_net_inflow, super_net_inflow_pct,
                   big_net_inflow, big_net_inflow_pct,
                   mid_net_inflow, mid_net_inflow_pct,
                   small_net_inflow, small_net_inflow_pct
            FROM stock_fund_flow_history 
            WHERE code = %s
            ORDER BY trade_date DESC
            LIMIT %s
        """, (code, limit))
        rows = cur.fetchall()
        
        # 获取最新日期
        latest_date = rows[0]['trade_date'] if rows else None
        
        return [dict(row) for row in rows], latest_date
    finally:
        put_db(conn)


def db_save_fund_flow_history(code, records):
    """保存资金流向历史数据到数据库"""
    conn = get_db()
    if not conn:
        return False
    
    try:
        cur = conn.cursor()
        for r in records:
            cur.execute("""
                INSERT INTO stock_fund_flow_history
                (code, trade_date, close_price, change_pct,
                 main_net_inflow, main_net_inflow_pct,
                 super_net_inflow, super_net_inflow_pct,
                 big_net_inflow, big_net_inflow_pct,
                 mid_net_inflow, mid_net_inflow_pct,
                 small_net_inflow, small_net_inflow_pct)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                ON CONFLICT (code, trade_date) DO UPDATE SET
                    close_price = EXCLUDED.close_price,
                    change_pct = EXCLUDED.change_pct,
                    main_net_inflow = EXCLUDED.main_net_inflow,
                    main_net_inflow_pct = EXCLUDED.main_net_inflow_pct,
                    super_net_inflow = EXCLUDED.super_net_inflow,
                    super_net_inflow_pct = EXCLUDED.super_net_inflow_pct,
                    big_net_inflow = EXCLUDED.big_net_inflow,
                    big_net_inflow_pct = EXCLUDED.big_net_inflow_pct,
                    mid_net_inflow = EXCLUDED.mid_net_inflow,
                    mid_net_inflow_pct = EXCLUDED.mid_net_inflow_pct,
                    small_net_inflow = EXCLUDED.small_net_inflow,
                    small_net_inflow_pct = EXCLUDED.small_net_inflow_pct,
                    updated_at = NOW()
            """, (
                code,
                r.get('日期') or r.get('trade_date'),
                r.get('收盘价') or r.get('close_price'),
                r.get('涨跌幅') or r.get('change_pct'),
                r.get('主力净流入-净额') or r.get('main_net_inflow'),
                r.get('主力净流入-净占比') or r.get('main_net_inflow_pct'),
                r.get('超大单净流入-净额') or r.get('super_net_inflow'),
                r.get('超大单净流入-净占比') or r.get('super_net_inflow_pct'),
                r.get('大单净流入-净额') or r.get('big_net_inflow'),
                r.get('大单净流入-净占比') or r.get('big_net_inflow_pct'),
                r.get('中单净流入-净额') or r.get('mid_net_inflow'),
                r.get('中单净流入-净占比') or r.get('mid_net_inflow_pct'),
                r.get('小单净流入-净额') or r.get('small_net_inflow'),
                r.get('小单净流入-净占比') or r.get('small_net_inflow_pct')
            ))
        conn.commit()
        return True
    except Exception as e:
        conn.rollback()
        print(f"保存资金流向历史失败: {e}")
        return False
    finally:
        put_db(conn)