#!/usr/bin/env python3
"""
全量股票技术信号扫描脚本
对数据库中的全部股票进行7个技术信号检测，结果存入 stock_signal_scan 表
支持断点续扫、并发处理、进度报告
"""

import sys
import os
import time
import json
import threading
import signal as sig_module
from datetime import datetime, date, timedelta
from concurrent.futures import ThreadPoolExecutor, as_completed

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import pandas as pd
import psycopg2
from psycopg2.extras import Json

from config import Config
from services.signal_detector import detect_all_signals
from services.stock_algorithms import (
    get_kline_from_local_db_threaded,
    get_ali_session, get_tencent_session,
    code_to_ali_symbol, code_to_tencent_symbol, is_bj_stock,
    ALICLOUD_KLINE_URL, TENCENT_KLINE_URL,
)

WORKERS = 8
BATCH_SIZE = 80
BATCH_SAVE_SIZE = 40
K_DAYS = 120
LOOKBACK = 5

_shutdown = False


def signal_handler(signum, frame):
    global _shutdown
    print("\n⚠️ 收到中断信号，正在优雅退出...")
    _shutdown = True


sig_module.signal(sig_module.SIGINT, signal_handler)
sig_module.signal(sig_module.SIGTERM, signal_handler)


def get_db_conn():
    return psycopg2.connect(
        host=Config.DB_HOST, port=Config.DB_PORT,
        dbname=Config.DB_NAME, user=Config.DB_USER, password=Config.DB_PASSWORD,
    )


def get_all_stock_codes(conn):
    """获取可交易的股票列表（排除退市、停牌、ST等无效股票）"""
    with conn.cursor() as cur:
        cur.execute("""
            SELECT code, name FROM stock_realtime_price
            WHERE volume > 0 AND price > 0
              AND name NOT LIKE '%%退%%'
              AND name NOT LIKE '%%ST%%'
              AND name NOT LIKE 'PT%%'
            ORDER BY code
        """)
        return cur.fetchall()


def get_scanned_codes(conn, scan_date):
    with conn.cursor() as cur:
        cur.execute(
            "SELECT code FROM stock_signal_scan WHERE scan_date = %s",
            (scan_date,),
        )
        return {row[0] for row in cur.fetchall()}


def get_kline_data(stock_code, days=K_DAYS):
    """获取K线数据：优先本地DB → 阿里云API → 腾讯API → 麦蕊API → AKShare
    （使用 services.stock_algorithms 统一的API会话和工具函数）"""
    from services.mairui_api import get_kline as _mairui_get_kline

    # 1. 优先从本地数据库读取（多线程安全版本）
    df = get_kline_from_local_db_threaded(stock_code, days)
    if df is not None:
        return df

    # 2. 阿里云K线API（北交所直接跳过，不支持）
    if not is_bj_stock(stock_code):
        try:
            session = get_ali_session()
            symbol = code_to_ali_symbol(stock_code)
            resp = session.post(ALICLOUD_KLINE_URL, data={
                'symbol': symbol, 'type': '240',
                'limit': str(min(days, 300)), 'ma': '5',
            }, timeout=10)
            if resp.status_code == 200:
                data = resp.json()
                if data.get('success') and data.get('data', {}).get('list'):
                    records = []
                    for item in data['data']['list']:
                        day_str = item.get('day', '')
                        if not day_str or len(day_str) < 10:
                            continue
                        records.append({
                            'date': day_str[:10],
                            'open': float(item.get('open', 0)),
                            'high': float(item.get('high', 0)),
                            'low': float(item.get('low', 0)),
                            'close': float(item.get('close', 0)),
                            'volume': float(item.get('volume', 0)),
                        })
                    if len(records) >= 30:
                        return pd.DataFrame(records)
        except Exception:
            pass

    # 3. 腾讯K线API（全市场，含北交所）
    try:
        session = get_tencent_session()
        tencent_symbol = code_to_tencent_symbol(stock_code)
        start_date = (datetime.now() - timedelta(days=days + 30)).strftime('%Y-%m-%d')
        resp = session.get(TENCENT_KLINE_URL, params={
            'param': f'{tencent_symbol},day,{start_date},,{min(days, 300)},qfq',
        }, timeout=15)
        if resp.status_code == 200:
            data = resp.json()
            stock_data = data.get('data', {}).get(tencent_symbol, {})
            klines = stock_data.get('qfqday') or stock_data.get('day') or []
            if len(klines) >= 30:
                records = []
                for item in klines:
                    if len(item) < 6:
                        continue
                    records.append({
                        'date': item[0][:10],
                        'open': float(item[1]),
                        'high': float(item[3]),
                        'low': float(item[4]),
                        'close': float(item[2]),
                        'volume': float(item[5]),
                    })
                if len(records) >= 30:
                    return pd.DataFrame(records)
    except Exception:
        pass

    # 4. 回退到麦蕊API
    try:
        result = _mairui_get_kline(stock_code, period='d', days=days, adjust='f')
        if result['success'] and result['data']:
            df = pd.DataFrame(result['data'])
            df.rename(columns={
                'date': 'date', 'open': 'open', 'high': 'high',
                'low': 'low', 'close': 'close', 'volume': 'volume',
            }, inplace=True)
            if len(df) >= 30:
                return df
    except Exception:
        pass

    # 5. 最后回退到AKShare (支持自动降级到腾讯数据源)
    try:
        from utils.data_fetcher import fetch_stock_hist
        end_date = datetime.now().strftime('%Y%m%d')
        start_date = (datetime.now() - timedelta(days=days)).strftime('%Y%m%d')
        df = fetch_stock_hist(
            stock_code=stock_code, period='daily',
            start_date=start_date, end_date=end_date, adjust='qfq',
        )
        if df is not None and not df.empty:
            df = df.rename(columns={
                '日期': 'date', '开盘': 'open', '最高': 'high',
                '最低': 'low', '收盘': 'close', '成交量': 'volume',
            })
            df = df[['date', 'open', 'high', 'low', 'close', 'volume']]
            return df
    except Exception:
        pass

    return None


def scan_single_stock(code, name):
    try:
        df = get_kline_data(code)
        if df is None or df.empty or len(df) < 30:
            return None

        result = detect_all_signals(df, lookback=LOOKBACK)
        signal_status = result.get('signal_status', [])
        triggered_count = sum(1 for s in signal_status if s.get('triggered'))

        return {
            'code': code,
            'name': name,
            'triggered_count': triggered_count,
            'signal_status': signal_status,
            'indicators': result.get('indicators', {}),
            'latest_signals': result.get('latest_signals', []),
        }
    except Exception:
        return None


def _fix_sequence(conn):
    """修复序列号，确保不会产生主键冲突"""
    try:
        with conn.cursor() as cur:
            cur.execute("""
                SELECT setval('stock_signal_scan_id_seq',
                    COALESCE((SELECT max(id) FROM stock_signal_scan), 0) + 1, false
                )
            """)
            conn.commit()
    except Exception as e:
        conn.rollback()
        print(f"⚠️ 修复序列失败: {e}", flush=True)


def save_batch(conn, results, scan_date):
    if not results:
        return

    # 每批次开始前确保序列正确
    _fix_sequence(conn)

    saved = 0
    for r in results:
        try:
            with conn.cursor() as cur:
                cur.execute("SAVEPOINT sp_insert")
                cur.execute("""
                    INSERT INTO stock_signal_scan
                        (code, name, scan_date, triggered_count, signal_status, indicators, latest_signals)
                    VALUES (%s, %s, %s, %s, %s, %s, %s)
                    ON CONFLICT (code, scan_date) DO UPDATE SET
                        name = EXCLUDED.name,
                        triggered_count = EXCLUDED.triggered_count,
                        signal_status = EXCLUDED.signal_status,
                        indicators = EXCLUDED.indicators,
                        latest_signals = EXCLUDED.latest_signals,
                        created_at = CURRENT_TIMESTAMP
                """, (
                    r['code'], r['name'], scan_date, r['triggered_count'],
                    Json(r['signal_status']), Json(r['indicators']), Json(r['latest_signals']),
                ))
                cur.execute("RELEASE SAVEPOINT sp_insert")
                saved += 1
        except Exception as e:
            # 主键冲突时回滚到 savepoint，修复序列后重试
            with conn.cursor() as cur:
                cur.execute("ROLLBACK TO SAVEPOINT sp_insert")
            if 'UniqueViolation' in type(e).__name__ or 'duplicate key' in str(e):
                _fix_sequence(conn)
                try:
                    with conn.cursor() as cur:
                        cur.execute("SAVEPOINT sp_insert")
                        cur.execute("""
                            INSERT INTO stock_signal_scan
                                (code, name, scan_date, triggered_count, signal_status, indicators, latest_signals)
                            VALUES (%s, %s, %s, %s, %s, %s, %s)
                            ON CONFLICT (code, scan_date) DO UPDATE SET
                                name = EXCLUDED.name,
                                triggered_count = EXCLUDED.triggered_count,
                                signal_status = EXCLUDED.signal_status,
                                indicators = EXCLUDED.indicators,
                                latest_signals = EXCLUDED.latest_signals,
                                created_at = CURRENT_TIMESTAMP
                        """, (
                            r['code'], r['name'], scan_date, r['triggered_count'],
                            Json(r['signal_status']), Json(r['indicators']), Json(r['latest_signals']),
                        ))
                        cur.execute("RELEASE SAVEPOINT sp_insert")
                        saved += 1
                except Exception as e2:
                    with conn.cursor() as cur:
                        cur.execute("ROLLBACK TO SAVEPOINT sp_insert")
                    print(f"⚠️ 重试保存失败 {r.get('code','?')}: {e2}", flush=True)
            else:
                print(f"⚠️ 保存失败 {r.get('code','?')}: {e}", flush=True)
    conn.commit()


def main():
    global _shutdown

    scan_date = date.today()

    force_rescan = os.environ.get('FORCE_RESCAN', '').strip() == '1'
    print(f"{'='*60}", flush=True)
    print(f"📊 全量股票技术信号扫描（高速版）", flush=True)
    print(f"📅 扫描日期: {scan_date}", flush=True)
    if force_rescan:
        print(f"⚠️  强制重新扫描模式", flush=True)
    print(f"⚙️  并发数: {WORKERS}, 批保存: {BATCH_SAVE_SIZE}, K线天数: {K_DAYS}", flush=True)
    print(f"{'='*60}", flush=True)

    conn = get_db_conn()

    if force_rescan:
        with conn.cursor() as cur:
            cur.execute("DELETE FROM stock_signal_scan WHERE scan_date = %s", (scan_date,))
            deleted = cur.rowcount
            conn.commit()
            print(f"🗑️  已清除今日 {deleted} 条扫描记录", flush=True)

    all_stocks = get_all_stock_codes(conn)
    total = len(all_stocks)
    print(f"📈 数据库股票总数: {total}", flush=True)

    scanned = get_scanned_codes(conn, scan_date)
    if scanned:
        print(f"✅ 今日已扫描: {len(scanned)} 只（续扫模式）", flush=True)

    # 过滤退市/ST股票 — 不参与扫描
    _SKIP_TAGS = ('退', 'ST', '*ST', '退市')
    pending = [(code, name) for code, name in all_stocks
               if code not in scanned and not any(tag in name for tag in _SKIP_TAGS)]
    skipped_st = len([1 for _, name in all_stocks if any(tag in name for tag in _SKIP_TAGS)])
    pending_count = len(pending)
    if skipped_st > 0:
        print(f"🚫 已过滤退市/ST股票: {skipped_st} 只", flush=True)
    print(f"⏳ 待扫描: {pending_count} 只", flush=True)

    if pending_count == 0:
        print("🎉 今日扫描已全部完成！", flush=True)
        show_summary(conn, scan_date)
        conn.close()
        return

    # 检测本地K线数据是否可用
    with conn.cursor() as cur:
        cur.execute("SELECT count(DISTINCT code) FROM stock_kline_daily WHERE trade_date >= CURRENT_DATE - INTERVAL '7 days'")
        local_kline_count = cur.fetchone()[0]
    if local_kline_count > 0:
        print(f"💾 本地K线数据: {local_kline_count} 只股票可用（优先使用本地数据）", flush=True)
    else:
        print(f"⚠️  本地无K线数据，将通过API获取（较慢）", flush=True)

    start_time = time.time()
    done_count = len(scanned)
    error_count = 0
    triggered_total = 0
    total_to_scan = total
    save_buffer = []
    last_report_time = time.time()

    print(f"\n🚀 开始扫描...", flush=True)
    print(f"-" * 60, flush=True)

    # 使用全局线程池 — 避免反复创建/销毁线程池开销
    with ThreadPoolExecutor(max_workers=WORKERS) as executor:
        futures = {}
        # 提交所有任务
        for code, name in pending:
            if _shutdown:
                break
            futures[executor.submit(scan_single_stock, code, name)] = (code, name)

        for future in as_completed(futures):
            if _shutdown:
                print("⏹️ 用户中断，正在保存当前进度...", flush=True)
                break

            code, name = futures[future]
            done_count += 1

            try:
                result = future.result()
            except Exception:
                result = None

            if result:
                save_buffer.append(result)
                if result['triggered_count'] > 0:
                    triggered_total += 1
                    names = [s['name'] for s in result['signal_status'] if s.get('triggered')]
                    print(f"  🔔 {result['code']} {result['name']}: {', '.join(names)}", flush=True)
            else:
                error_count += 1

            # 攒够一批就保存（减少DB写入频率）
            if len(save_buffer) >= BATCH_SAVE_SIZE:
                save_batch(conn, save_buffer, scan_date)
                save_buffer = []

            # 每3秒报告一次进度（避免刷屏）
            now = time.time()
            if now - last_report_time >= 3:
                elapsed = now - start_time
                scanned_this_run = done_count - len(scanned)
                speed = scanned_this_run / elapsed if elapsed > 0 else 0
                remaining_stocks = total_to_scan - done_count
                remaining_time = remaining_stocks / speed if speed > 0 else 0
                pct = done_count / total_to_scan * 100
                print(f"  [{pct:5.1f}%] {done_count}/{total_to_scan} "
                      f"| 速度: {speed:.1f}只/秒 | 剩余: {remaining_time/60:.1f}分钟 "
                      f"| 触发: {triggered_total} | 失败: {error_count}", flush=True)
                last_report_time = now

    # 保存剩余结果
    if save_buffer:
        save_batch(conn, save_buffer, scan_date)

    elapsed = time.time() - start_time
    print(f"\n{'='*60}", flush=True)
    print(f"✅ 扫描{'中断' if _shutdown else '完成'}！", flush=True)
    print(f"  扫描: {done_count} 只 | 耗时: {elapsed/60:.1f}分钟", flush=True)
    print(f"  触发信号: {triggered_total} 只 | 失败: {error_count} 只", flush=True)
    final_speed = (done_count - len(scanned)) / elapsed if elapsed > 0 else 0
    print(f"  平均速度: {final_speed:.1f} 只/秒", flush=True)
    print(f"{'='*60}", flush=True)

    show_summary(conn, scan_date)
    conn.close()


def show_summary(conn, scan_date):
    print(f"\n📊 扫描结果摘要（{scan_date}）", flush=True)
    print(f"-" * 60, flush=True)

    with conn.cursor() as cur:
        cur.execute("""
            SELECT count(*),
                   coalesce(sum(case when triggered_count > 0 then 1 else 0 end), 0)
            FROM stock_signal_scan WHERE scan_date = %s
        """, (scan_date,))
        total, triggered = cur.fetchone()
        print(f"  总扫描: {total} 只 | 有信号: {triggered} 只", flush=True)

        cur.execute("""
            SELECT code, name, triggered_count, signal_status
            FROM stock_signal_scan
            WHERE scan_date = %s AND triggered_count > 0
            ORDER BY triggered_count DESC
            LIMIT 30
        """, (scan_date,))
        rows = cur.fetchall()

        if rows:
            print(f"\n🔔 触发信号TOP30：", flush=True)
            for code, name, tc, status in rows:
                signals = status if isinstance(status, list) else json.loads(status) if status else []
                names = [s['name'] for s in signals if s.get('triggered')]
                print(f"  {code} {name:8s} | {tc}个信号: {', '.join(names)}", flush=True)
        else:
            print(f"\n  暂无触发信号的股票", flush=True)

        cur.execute("""
            SELECT
                s.value->>'name' as signal_name,
                count(*) as cnt
            FROM stock_signal_scan, jsonb_array_elements(signal_status) s
            WHERE scan_date = %s AND (s.value->>'triggered')::boolean = true
            GROUP BY s.value->>'name'
            ORDER BY cnt DESC
        """, (scan_date,))
        signal_dist = cur.fetchall()
        if signal_dist:
            print(f"\n📈 信号分布：", flush=True)
            for name, cnt in signal_dist:
                print(f"  {name}: {cnt} 只", flush=True)


if __name__ == '__main__':
    main()