Stock Research Group

# analyzers/balancesheet/_shared/db.py
"""资产负债表数据库工具"""

import json
import sqlite3
from datetime import datetime
from pathlib import Path

from config import DB_PATH

SCHEMA_PATH = Path(__file__).parent.parent / "schema.sql"


def connect(db_path: Path = DB_PATH) -> sqlite3.Connection:
    """连接数据库"""
    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    return conn


def init_table(conn: sqlite3.Connection = None):
    """初始化表结构"""
    if conn is None:
        conn = connect()
        should_close = True
    else:
        should_close = False
    
    try:
        schema = SCHEMA_PATH.read_text(encoding="utf-8")
        conn.executescript(schema)
        conn.commit()
    finally:
        if should_close:
            conn.close()


def upsert_balancesheet(conn: sqlite3.Connection, rows: list) -> int:
    """批量插入/更新资产负债表数据
    
    Args:
        conn: 数据库连接
        rows: 数据行列表，每行包含所有字段
        
    Returns:
        插入/更新的行数
    """
    if not rows:
        return 0
    
    # 基础列（常用字段）
    base_cols = [
        'ts_code', 'ann_date', 'f_ann_date', 'end_date', 'report_type', 
        'comp_type', 'end_type', 'total_share', 'money_cap', 
        'accounts_receiv', 'inventories', 'total_cur_assets', 'fix_assets',
        'total_assets', 'st_borr', 'acct_payable', 'total_cur_liab',
        'total_liab', 'undistr_porfit', 'total_hldr_eqy_exc_min_int'
    ]
    
    # 唯一键：如果ann_date为空，使用ts_code+end_date+report_type作为唯一键
    # 这样可以确保同一报告期的数据能正确更新
    unique_cols = ['ts_code', 'end_date', 'report_type']
    
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # 准备数据
    db_rows = []
    for row in rows:
        db_row = {}
        # 提取基础列
        for col in base_cols:
            db_row[col] = row.get(col)
        
        # 存储完整JSON
        db_row['payload_json'] = json.dumps(row, ensure_ascii=False, default=str)
        db_row['source'] = 'tushare'
        db_row['created_at'] = now
        db_row['updated_at'] = now
        
        # 处理唯一键中的NULL值
        for col in unique_cols:
            if db_row.get(col) is None:
                db_row[col] = ''
        
        # 如果ann_date为空，设置为空字符串（不影响唯一键判断）
        if db_row.get('ann_date') is None:
            db_row['ann_date'] = ''
        
        db_rows.append(db_row)
    
    # 构建SQL
    all_cols = list(db_rows[0].keys())
    placeholders = ", ".join(["?"] * len(all_cols))
    columns_sql = ", ".join(all_cols)
    conflict_sql = ", ".join(unique_cols)
    update_cols = [c for c in all_cols if c not in unique_cols + ['created_at']]
    update_sql = ", ".join([f"{c} = excluded.{c}" for c in update_cols])
    
    sql = (
        f"INSERT INTO balancesheet ({columns_sql}) "
        f"VALUES ({placeholders}) "
        f"ON CONFLICT({conflict_sql}) DO UPDATE SET {update_sql}"
    )
    
    values = [[r.get(col) for col in all_cols] for r in db_rows]
    
    try:
        conn.executemany(sql, values)
        conn.commit()
        return len(db_rows)
    except Exception as e:
        conn.rollback()
        raise


if __name__ == "__main__":
    # CLI入口：初始化表
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--init", action="store_true", help="初始化表结构")
    args = parser.parse_args()
    
    if args.init:
        init_table()
        print("表结构初始化完成")

# analyzers/balancesheet/pipeline.py
"""资产负债表分析流程入口"""

import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(PROJECT_ROOT))

from analyzers.balancesheet.search import (
    get_balancesheet,
    get_balancesheet_history,
    get_field_value,
    search_by_field,
    ensure_data
)
from fetchers.balancesheet import fetch_and_save


def run_fetch(**kwargs):
    """拉取数据"""
    count = fetch_and_save(**kwargs)
    print(f"保存 {count} 条记录")
    return count


def run_query(ts_code: str, field: str = None, end_date: str = None):
    """查询数据"""
    if field:
        value = get_field_value(ts_code, field, end_date)
        if value is not None:
            print(f"{ts_code} {field} = {value:,.0f}")
        else:
            print(f"未找到数据")
    else:
        record = get_balancesheet(ts_code, end_date)
        if record:
            print(f"报告期: {record.get('end_date')}")
            total_assets = record.get('total_assets') or 0
            inventories = record.get('inventories') or 0
            accounts_receiv = record.get('accounts_receiv') or 0
            print(f"总资产: {total_assets:,.0f}")
            print(f"存货: {inventories:,.0f}")
            print(f"应收账款: {accounts_receiv:,.0f}")
        else:
            print("未找到数据")


def run_history(ts_code: str, limit: int = 4):
    """查询历史数据"""
    records = get_balancesheet_history(ts_code, limit=limit)
    print(f"\n{ts_code} 最近{len(records)}期数据:")
    print("-" * 60)
    for r in records:
        total_assets = r.get('total_assets') or 0
        inventories = r.get('inventories') or 0
        print(f"{r['end_date']}: 总资产={total_assets:,.0f}, "
              f"存货={inventories:,.0f}")


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="资产负债表分析")
    subparsers = parser.add_subparsers(dest="command")
    
    # fetch命令
    p_fetch = subparsers.add_parser("fetch", help="拉取数据")
    p_fetch.add_argument("--ts-code")
    p_fetch.add_argument("--start-date")
    p_fetch.add_argument("--end-date")
    p_fetch.add_argument("--period", help="报告期 YYYYMMDD（VIP接口必需）")
    p_fetch.add_argument("--report-type", default="1")
    p_fetch.add_argument("--comp-type", help="公司类型")
    p_fetch.add_argument("--vip", action="store_true", help="使用VIP接口（全量拉取，需提供period）")
    
    # query命令
    p_query = subparsers.add_parser("query", help="查询数据")
    p_query.add_argument("--ts-code", required=True)
    p_query.add_argument("--field")
    p_query.add_argument("--end-date")
    
    # history命令
    p_history = subparsers.add_parser("history", help="查询历史")
    p_history.add_argument("--ts-code", required=True)
    p_history.add_argument("--limit", type=int, default=4)
    
    args = parser.parse_args()
    
    if args.command == "fetch":
        if args.vip and not args.period:
            parser.error("使用VIP接口时必须提供 --period 参数")
        run_fetch(
            ts_code=args.ts_code,
            start_date=args.start_date,
            end_date=args.end_date,
            period=args.period,
            report_type=args.report_type,
            comp_type=args.comp_type,
            use_vip=args.vip
        )
    elif args.command == "query":
        run_query(args.ts_code, args.field, args.end_date)
    elif args.command == "history":
        run_history(args.ts_code, args.limit)

PRAGMA journal_mode=WAL;

CREATE TABLE IF NOT EXISTS balancesheet (
  ts_code TEXT NOT NULL,
  ann_date TEXT,
  f_ann_date TEXT,
  end_date TEXT NOT NULL,
  report_type TEXT,
  comp_type TEXT,
  end_type TEXT,
  total_share REAL,
  money_cap REAL,
  accounts_receiv REAL,
  inventories REAL,
  total_cur_assets REAL,
  fix_assets REAL,
  total_assets REAL,
  st_borr REAL,
  acct_payable REAL,
  total_cur_liab REAL,
  total_liab REAL,
  undistr_porfit REAL,
  total_hldr_eqy_exc_min_int REAL,
  source TEXT,
  payload_json TEXT NOT NULL,
  created_at TEXT NOT NULL,
  updated_at TEXT NOT NULL,
  UNIQUE(ts_code, end_date, report_type)
);

CREATE INDEX IF NOT EXISTS idx_balancesheet_ts_code ON balancesheet(ts_code);
CREATE INDEX IF NOT EXISTS idx_balancesheet_end_date ON balancesheet(end_date);
CREATE INDEX IF NOT EXISTS idx_balancesheet_ts_end ON balancesheet(ts_code, end_date);

# analyzers/balancesheet/search.py
"""资产负债表查询功能"""

import json
import sys
from pathlib import Path
from typing import Optional, List, Dict

PROJECT_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(PROJECT_ROOT))

from analyzers.balancesheet._shared.db import connect


def get_balancesheet(
    ts_code: str,
    end_date: str = None,
    report_type: str = "1"
) -> Optional[Dict]:
    """
    查询单条资产负债表记录
    
    Args:
        ts_code: 股票代码
        end_date: 报告期 YYYYMMDD，默认最新
        report_type: 报表类型，默认1（合并报表）
        
    Returns:
        记录字典，包含基础列和完整字段
    """
    conn = connect()
    try:
        if end_date:
            sql = """
                SELECT * FROM balancesheet 
                WHERE ts_code = ? AND end_date = ? AND report_type = ?
                ORDER BY ann_date DESC
                LIMIT 1
            """
            cur = conn.execute(sql, (ts_code, end_date, report_type))
        else:
            sql = """
                SELECT * FROM balancesheet 
                WHERE ts_code = ? AND report_type = ?
                ORDER BY end_date DESC, ann_date DESC
                LIMIT 1
            """
            cur = conn.execute(sql, (ts_code, report_type))
        
        row = cur.fetchone()
        if not row:
            return None
        
        result = dict(row)
        # 解析完整字段
        if result.get('payload_json'):
            full_data = json.loads(result['payload_json'])
            result.update(full_data)
        
        return result
    finally:
        conn.close()


def get_balancesheet_history(
    ts_code: str,
    start_date: str = None,
    end_date: str = None,
    report_type: str = "1",
    limit: int = None
) -> List[Dict]:
    """
    查询历史资产负债表记录
    
    Args:
        ts_code: 股票代码
        start_date: 开始日期 YYYYMMDD
        end_date: 结束日期 YYYYMMDD
        report_type: 报表类型
        limit: 限制返回数量（如4表示最近4个季度）
        
    Returns:
        记录列表，按end_date降序
    """
    conn = connect()
    try:
        conditions = ["ts_code = ?", "report_type = ?"]
        params = [ts_code, report_type]
        
        if start_date:
            conditions.append("end_date >= ?")
            params.append(start_date)
        if end_date:
            conditions.append("end_date <= ?")
            params.append(end_date)
        
        sql = f"""
            SELECT * FROM balancesheet 
            WHERE {' AND '.join(conditions)}
            ORDER BY end_date DESC, ann_date DESC
        """
        
        if limit:
            sql += f" LIMIT {limit}"
        
        cur = conn.execute(sql, params)
        rows = cur.fetchall()
        
        results = []
        for row in rows:
            result = dict(row)
            # 解析完整字段
            if result.get('payload_json'):
                full_data = json.loads(result['payload_json'])
                result.update(full_data)
            results.append(result)
        
        return results
    finally:
        conn.close()


def ensure_data(
    ts_code: str,
    end_date: str = None,
    years: int = 1,
    report_type: str = "1"
) -> bool:
    """
    确保有足够的历史数据，不足时自动拉取
    
    Args:
        ts_code: 股票代码
        end_date: 截止日期，默认当前日期
        years: 需要的数据年数
        report_type: 报表类型
        
    Returns:
        是否成功获取数据
    """
    from datetime import datetime, timedelta
    from fetchers.balancesheet import fetch_and_save
    
    if end_date is None:
        end_date = datetime.now().strftime("%Y%m%d")
    
    # 检查现有数据
    existing = get_balancesheet_history(
        ts_code=ts_code,
        end_date=end_date,
        report_type=report_type,
        limit=years * 4 + 4  # 多查一些确保
    )
    
    # 计算需要的日期范围（N+1年，确保有足够数据计算同比）
    need_years = years + 1
    start_date_dt = datetime.strptime(end_date, "%Y%m%d") - timedelta(days=need_years * 365)
    start_date = start_date_dt.strftime("%Y%m%d")
    
    # 如果数据不足，拉取
    if len(existing) < years * 4:
        print(f"数据不足（{len(existing)}条），拉取 {start_date} 至 {end_date} 的数据...")
        try:
            count = fetch_and_save(
                ts_code=ts_code,
                start_date=start_date,
                end_date=end_date,
                report_type=report_type
            )
            print(f"拉取完成，新增 {count} 条记录")
            return True
        except Exception as e:
            print(f"拉取失败: {e}")
            return False
    
    return True


def get_field_value(
    ts_code: str,
    field_name: str,
    end_date: str = None,
    report_type: str = "1",
    auto_fetch: bool = True
) -> Optional[float]:
    """
    查询特定字段值，支持自动拉取
    
    Args:
        ts_code: 股票代码
        field_name: 字段名（如'inventories'）
        end_date: 报告期，默认最新
        report_type: 报表类型
        auto_fetch: 是否自动拉取数据（如果数据库无数据）
        
    Returns:
        字段值（float），不存在返回None
    """
    if auto_fetch:
        ensure_data(ts_code, end_date, years=1, report_type=report_type)
    
    record = get_balancesheet(ts_code, end_date, report_type)
    if not record:
        return None
    
    # 先查基础列
    if field_name in record:
        value = record[field_name]
        if value is not None:
            return float(value)
    
    # 再从payload_json中查找
    if 'payload_json' in record:
        full_data = json.loads(record['payload_json'])
        value = full_data.get(field_name)
        if value is not None:
            return float(value)
    
    return None


def search_by_field(
    field_name: str,
    end_date: str,
    report_type: str = "1",
    comp_type: str = None,
    limit: int = 100,
    order: str = "DESC"
) -> List[Dict]:
    """
    跨公司查询特定字段
    
    Args:
        field_name: 字段名
        end_date: 报告期
        report_type: 报表类型
        comp_type: 公司类型过滤
        limit: 返回数量限制
        order: 排序方向（DESC/ASC）
        
    Returns:
        列表，每项包含ts_code和字段值
    """
    conn = connect()
    try:
        # 先检查是否是基础列
        base_cols = [
            'total_share', 'money_cap', 'accounts_receiv', 'inventories',
            'total_cur_assets', 'fix_assets', 'total_assets', 'st_borr',
            'acct_payable', 'total_cur_liab', 'total_liab', 'undistr_porfit',
            'total_hldr_eqy_exc_min_int'
        ]
        
        conditions = ["end_date = ?", "report_type = ?"]
        params = [end_date, report_type]
        
        if comp_type:
            conditions.append("comp_type = ?")
            params.append(comp_type)
        
        if field_name in base_cols:
            # 直接从基础列查询
            sql = f"""
                SELECT ts_code, {field_name} as value
                FROM balancesheet
                WHERE {' AND '.join(conditions)} AND {field_name} IS NOT NULL
                ORDER BY {field_name} {order}
                LIMIT ?
            """
            params.append(limit)
            cur = conn.execute(sql, params)
        else:
            # 需要解析JSON（性能较差，但支持所有字段）
            sql = f"""
                SELECT ts_code, payload_json
                FROM balancesheet
                WHERE {' AND '.join(conditions)}
                LIMIT ?
            """
            params.append(limit * 2)  # 多查一些，因为可能有些记录没有该字段
            cur = conn.execute(sql, params)
            
            # 解析并过滤
            results = []
            for row in cur.fetchall():
                try:
                    data = json.loads(row['payload_json'])
                    value = data.get(field_name)
                    if value is not None:
                        results.append({
                            'ts_code': row['ts_code'],
                            'value': float(value)
                        })
                except:
                    continue
            
            # 排序
            results.sort(key=lambda x: x['value'], reverse=(order == "DESC"))
            return results[:limit]
        
        return [dict(row) for row in cur.fetchall()]
    finally:
        conn.close()

"""券商金股数据库工具"""

import sqlite3
from pathlib import Path

from config import DB_PATH

SCHEMA_PATH = Path(__file__).parent.parent / "schema.sql"


def connect() -> sqlite3.Connection:
    conn = sqlite3.connect(str(DB_PATH))
    conn.row_factory = sqlite3.Row
    return conn


def init_table():
    """初始化券商金股表"""
    conn = connect()
    try:
        schema = SCHEMA_PATH.read_text(encoding="utf-8")
        conn.executescript(schema)
        conn.commit()
    finally:
        conn.close()


def upsert_recommendations(rows: list[dict]) -> int:
    """批量插入金股数据，返回插入条数"""
    if not rows:
        return 0
    conn = connect()
    try:
        sql = """
            INSERT INTO broker_recommend (month, broker, ts_code, name)
            VALUES (?, ?, ?, ?)
            ON CONFLICT(month, broker, ts_code) DO UPDATE SET
                name = excluded.name
        """
        values = [
            (r.get("month"), r.get("broker"), r.get("ts_code"), r.get("name"))
            for r in rows
        ]
        conn.executemany(sql, values)
        conn.commit()
        return len(rows)
    finally:
        conn.close()


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="券商金股数据库工具")
    parser.add_argument("--init", action="store_true", help="初始化表")
    args = parser.parse_args()
    
    if args.init:
        init_table()
        print(f"表已初始化: {DB_PATH}")
    else:
        parser.print_help()

# analyzers/broker_recommend/pipeline.py
"""券商金股分析流程入口"""

import json
import sys
from datetime import datetime
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(PROJECT_ROOT))

from fetchers.broker_recommend import fetch_broker_recommend
from analyzers.broker_recommend.stats import get_monthly_stats


def run_fetch(month: str = None) -> int:
    """
    拉取数据
    
    Args:
        month: 月度 YYYYMM，None 时使用当前月份
    
    Returns:
        新增条数
    """
    if not month:
        month = datetime.now().strftime("%Y%m")
    return fetch_broker_recommend(month)


def run_stats(month: str = None) -> dict:
    """
    运行统计
    
    Args:
        month: 月度 YYYYMM，None 时使用当前月份
    
    Returns:
        统计结果
    """
    if not month:
        month = datetime.now().strftime("%Y%m")
    return get_monthly_stats(month)


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="券商金股分析流程")
    subparsers = parser.add_subparsers(dest="command")
    
    # 拉取
    fetch_p = subparsers.add_parser("fetch", help="拉取金股数据")
    fetch_p.add_argument("--month", help="月度 YYYYMM，默认当前月份")
    
    # 统计
    stats_p = subparsers.add_parser("stats", help="查看统计")
    stats_p.add_argument("--month", help="月度 YYYYMM，默认当前月份")
    
    args = parser.parse_args()
    
    if args.command == "fetch":
        count = run_fetch(args.month)
        print(f"已入库 {count} 条")
    
    elif args.command == "stats":
        result = run_stats(args.month)
        print(json.dumps(result, ensure_ascii=False, indent=2))
    
    else:
        parser.print_help()

-- analyzers/broker_recommend/schema.sql
CREATE TABLE IF NOT EXISTS broker_recommend (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    month TEXT NOT NULL,              -- 月度 YYYYMM
    broker TEXT NOT NULL,             -- 券商名称
    ts_code TEXT NOT NULL,            -- 股票代码
    name TEXT NOT NULL,               -- 股票简称
    created_at TEXT DEFAULT CURRENT_TIMESTAMP,
    UNIQUE(month, broker, ts_code)
);

CREATE INDEX IF NOT EXISTS idx_br_month ON broker_recommend(month);
CREATE INDEX IF NOT EXISTS idx_br_ts_code ON broker_recommend(ts_code);
CREATE INDEX IF NOT EXISTS idx_br_broker ON broker_recommend(broker);

# analyzers/broker_recommend/stats.py
"""券商金股统计功能"""

from collections import defaultdict
from ._shared.db import connect


def get_monthly_stats(month: str) -> dict:
    """
    获取指定月份的统计
    
    Returns:
        {
            "month": "202602",
            "by_company": [{"ts_code": "...", "name": "...", "count": 5, "brokers": [...]}],
            "by_broker": [{"broker": "...", "count": 10}]
        }
    """
    conn = connect()
    try:
        cur = conn.execute(
            "SELECT ts_code, name, broker FROM broker_recommend WHERE month = ?",
            (month,)
        )
        rows = [dict(r) for r in cur.fetchall()]
    finally:
        conn.close()
    
    # 按公司统计
    company_stats = defaultdict(lambda: {"ts_code": "", "name": "", "count": 0, "brokers": []})
    for row in rows:
        key = row["ts_code"]
        company_stats[key]["ts_code"] = row["ts_code"]
        company_stats[key]["name"] = row["name"]
        company_stats[key]["count"] += 1
        company_stats[key]["brokers"].append(row["broker"])
    
    by_company = sorted(
        [
            {
                "ts_code": v["ts_code"],
                "name": v["name"],
                "count": v["count"],
                "brokers": list(set(v["brokers"]))
            }
            for v in company_stats.values()
        ],
        key=lambda x: x["count"],
        reverse=True
    )
    
    # 按券商统计
    broker_stats = defaultdict(int)
    for row in rows:
        broker_stats[row["broker"]] += 1
    
    by_broker = sorted(
        [{"broker": k, "count": v} for k, v in broker_stats.items()],
        key=lambda x: x["count"],
        reverse=True
    )
    
    return {
        "month": month,
        "by_company": by_company,
        "by_broker": by_broker
    }


def get_continuous_recommendations(ts_code: str, min_months: int = 2) -> list[str]:
    """
    获取连续被推荐的月份列表
    
    Args:
        ts_code: 股票代码
        min_months: 最少连续月份数
    
    Returns:
        连续月份列表，如 ["202501", "202502", "202503"]
    """
    conn = connect()
    try:
        cur = conn.execute(
            "SELECT DISTINCT month FROM broker_recommend WHERE ts_code = ? ORDER BY month",
            (ts_code,)
        )
        months = [r["month"] for r in cur.fetchall()]
    finally:
        conn.close()
    
    if len(months) < min_months:
        return []
    
    # 找出连续月份（简化版：按字符串排序后检查是否连续）
    continuous = []
    current_seq = [months[0]]
    
    for i in range(1, len(months)):
        prev = int(months[i-1])
        curr = int(months[i])
        
        # 计算月份差
        prev_year = prev // 100
        prev_month = prev % 100
        curr_year = curr // 100
        curr_month = curr % 100
        
        # 判断是否连续
        if curr_year == prev_year and curr_month == prev_month + 1:
            # 同一年，月份+1
            current_seq.append(months[i])
        elif curr_year == prev_year + 1 and curr_month == 1 and prev_month == 12:
            # 跨年：12月 -> 1月
            current_seq.append(months[i])
        else:
            # 不连续
            if len(current_seq) >= min_months:
                continuous.append(current_seq)
            current_seq = [months[i]]
    
    if len(current_seq) >= min_months:
        continuous.append(current_seq)
    
    # 返回最长的连续序列
    return max(continuous, key=len) if continuous else []


def get_broker_recommendations(broker: str, month: str = None) -> list[dict]:
    """
    获取某券商的金股列表
    
    Args:
        broker: 券商名称
        month: 可选，指定月份
    
    Returns:
        金股列表
    """
    conn = connect()
    try:
        if month:
            cur = conn.execute(
                "SELECT month, ts_code, name FROM broker_recommend WHERE broker = ? AND month = ? ORDER BY ts_code",
                (broker, month)
            )
        else:
            cur = conn.execute(
                "SELECT month, ts_code, name FROM broker_recommend WHERE broker = ? ORDER BY month DESC, ts_code",
                (broker,)
            )
        return [dict(r) for r in cur.fetchall()]
    finally:
        conn.close()

"""现金流量表数据库工具"""

import json
import sqlite3
from datetime import datetime
from pathlib import Path

from config import DB_PATH

SCHEMA_PATH = Path(__file__).parent.parent / "schema.sql"


def connect(db_path: Path = DB_PATH) -> sqlite3.Connection:
    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    return conn


def init_table(conn: sqlite3.Connection = None):
    if conn is None:
        conn = connect()
        should_close = True
    else:
        should_close = False
    try:
        schema = SCHEMA_PATH.read_text(encoding="utf-8")
        conn.executescript(schema)
        conn.commit()
    finally:
        if should_close:
            conn.close()


def upsert_cashflow(conn: sqlite3.Connection, rows: list) -> int:
    if not rows:
        return 0

    base_cols = [
        "ts_code", "ann_date", "end_date",
        "n_cashflow_act", "n_cashflow_inv", "n_cashflow_fin",
        "c_cash_equ_beg_period", "c_cash_equ_end_period",
        "c_inf_fr_operate_a", "c_outf_operate_a",
    ]

    unique_cols = ["ts_code", "end_date"]
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    db_rows = []
    for row in rows:
        db_row = {col: row.get(col) for col in base_cols}
        if not db_row.get("ts_code") or not db_row.get("end_date"):
            continue
        db_row["payload_json"] = json.dumps(row, ensure_ascii=False, default=str)
        db_row["source"] = "tushare"
        db_row["created_at"] = now
        db_row["updated_at"] = now
        db_rows.append(db_row)

    if not db_rows:
        return 0

    all_cols = list(db_rows[0].keys())
    placeholders = ", ".join(["?"] * len(all_cols))
    columns_sql = ", ".join(all_cols)
    conflict_sql = ", ".join(unique_cols)
    update_cols = [c for c in all_cols if c not in unique_cols + ["created_at"]]
    update_sql = ", ".join([f"{c} = excluded.{c}" for c in update_cols])

    sql = (
        f"INSERT INTO cashflow ({columns_sql}) "
        f"VALUES ({placeholders}) "
        f"ON CONFLICT({conflict_sql}) DO UPDATE SET {update_sql}"
    )
    values = [[r.get(col) for col in all_cols] for r in db_rows]

    try:
        conn.executemany(sql, values)
        conn.commit()
        return len(db_rows)
    except Exception:
        conn.rollback()
        raise


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--init", action="store_true", help="初始化表结构")
    args = parser.parse_args()
    if args.init:
        init_table()
        print("表结构初始化完成")

"""现金流量表流程入口"""

import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(PROJECT_ROOT))

from analyzers.cashflow.search import get_cashflow, get_cashflow_history, get_field_value, ensure_data
from fetchers.cashflow import fetch_and_save


def run_fetch(**kwargs):
    count = fetch_and_save(**kwargs)
    print(f"保存 {count} 条记录")
    return count


def run_query(ts_code: str, field: str = None, end_date: str = None):
    ensure_data(ts_code, end_date, years=4)
    if field:
        value = get_field_value(ts_code, field, end_date)
        print("未找到数据" if value is None else f"{ts_code} {field} = {value}")
    else:
        record = get_cashflow(ts_code, end_date)
        print("未找到数据" if not record else record)


def run_history(ts_code: str, limit: int = 4):
    records = get_cashflow_history(ts_code, limit=limit)
    for r in records:
        print(r.get("end_date"), r.get("n_cashflow_act"), r.get("n_cashflow_inv"))


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="现金流量表查询")
    subparsers = parser.add_subparsers(dest="command", required=True)

    p_fetch = subparsers.add_parser("fetch", help="拉取数据")
    p_fetch.add_argument("--ts-code")
    p_fetch.add_argument("--ann-date")
    p_fetch.add_argument("--f-ann-date")
    p_fetch.add_argument("--start-date")
    p_fetch.add_argument("--end-date")
    p_fetch.add_argument("--period")
    p_fetch.add_argument("--report-type")
    p_fetch.add_argument("--comp-type")
    p_fetch.add_argument("--is-calc", type=int)
    p_fetch.add_argument("--vip", action="store_true")

    p_query = subparsers.add_parser("query", help="查询数据")
    p_query.add_argument("--ts-code", required=True)
    p_query.add_argument("--field")
    p_query.add_argument("--end-date")

    p_history = subparsers.add_parser("history", help="查询历史")
    p_history.add_argument("--ts-code", required=True)
    p_history.add_argument("--limit", type=int, default=4)

    args = parser.parse_args()
    if args.command == "fetch":
        if not any([args.ts_code, args.period, args.start_date, args.end_date]):
            parser.error("请至少提供一个查询条件：--ts-code / --period / --start-date / --end-date")
        if args.vip and not args.period:
            parser.error("使用VIP接口时必须提供 --period 参数")
        run_fetch(
            ts_code=args.ts_code,
            ann_date=args.ann_date,
            f_ann_date=args.f_ann_date,
            start_date=args.start_date,
            end_date=args.end_date,
            period=args.period,
            report_type=args.report_type,
            comp_type=args.comp_type,
            is_calc=args.is_calc,
            use_vip=args.vip,
        )
    elif args.command == "query":
        run_query(args.ts_code, args.field, args.end_date)
    elif args.command == "history":
        run_history(args.ts_code, args.limit)

PRAGMA journal_mode=WAL;

CREATE TABLE IF NOT EXISTS cashflow (
  ts_code TEXT NOT NULL,
  ann_date TEXT,
  end_date TEXT NOT NULL,
  n_cashflow_act REAL,
  n_cashflow_inv REAL,
  n_cashflow_fin REAL,
  c_cash_equ_beg_period REAL,
  c_cash_equ_end_period REAL,
  c_inf_fr_operate_a REAL,
  c_outf_operate_a REAL,
  source TEXT,
  payload_json TEXT NOT NULL,
  created_at TEXT NOT NULL,
  updated_at TEXT NOT NULL,
  UNIQUE(ts_code, end_date)
);

CREATE INDEX IF NOT EXISTS idx_cashflow_ts_code ON cashflow(ts_code);
CREATE INDEX IF NOT EXISTS idx_cashflow_end_date ON cashflow(end_date);
CREATE INDEX IF NOT EXISTS idx_cashflow_ts_end ON cashflow(ts_code, end_date);

"""现金流量表查询功能"""

import json
import sys
from pathlib import Path
from typing import Optional, List, Dict

PROJECT_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(PROJECT_ROOT))

from analyzers.cashflow._shared.db import connect


def get_cashflow(ts_code: str, end_date: str = None) -> Optional[Dict]:
    conn = connect()
    try:
        if end_date:
            sql = "SELECT * FROM cashflow WHERE ts_code = ? AND end_date = ? ORDER BY ann_date DESC LIMIT 1"
            cur = conn.execute(sql, (ts_code, end_date))
        else:
            sql = "SELECT * FROM cashflow WHERE ts_code = ? ORDER BY end_date DESC, ann_date DESC LIMIT 1"
            cur = conn.execute(sql, (ts_code,))
        row = cur.fetchone()
        if not row:
            return None
        result = dict(row)
        if result.get("payload_json"):
            try:
                result.update(json.loads(result["payload_json"]))
            except (json.JSONDecodeError, TypeError):
                pass
        return result
    finally:
        conn.close()


def get_cashflow_history(
    ts_code: str,
    start_date: str = None,
    end_date: str = None,
    limit: int = None
) -> List[Dict]:
    conn = connect()
    try:
        conditions = ["ts_code = ?"]
        params = [ts_code]
        if start_date:
            conditions.append("end_date >= ?")
            params.append(start_date)
        if end_date:
            conditions.append("end_date <= ?")
            params.append(end_date)
        sql = f"SELECT * FROM cashflow WHERE {' AND '.join(conditions)} ORDER BY end_date DESC, ann_date DESC"
        if limit is not None:
            if not isinstance(limit, int) or limit <= 0:
                raise ValueError("limit must be a positive int")
            sql += " LIMIT ?"
            params.append(limit)
        cur = conn.execute(sql, params)
        rows = cur.fetchall()
        results = []
        for row in rows:
            result = dict(row)
            if result.get("payload_json"):
                try:
                    result.update(json.loads(result["payload_json"]))
                except (json.JSONDecodeError, TypeError):
                    pass
            results.append(result)
        return results
    finally:
        conn.close()


def get_field_value(ts_code: str, field_name: str, end_date: str = None) -> Optional[float]:
    record = get_cashflow(ts_code, end_date)
    if not record:
        return None
    value = record.get(field_name)
    if value is not None:
        return float(value)
    return None


def ensure_data(ts_code: str, end_date: str = None, years: int = 4) -> bool:
    from datetime import datetime, timedelta
    from fetchers.cashflow import fetch_and_save

    if end_date is None:
        end_date = datetime.now().strftime("%Y%m%d")
    existing = get_cashflow_history(ts_code, end_date=end_date, limit=years * 4 + 4)
    need_years = years + 1
    start_date_dt = datetime.strptime(end_date, "%Y%m%d") - timedelta(days=need_years * 365)
    start_date = start_date_dt.strftime("%Y%m%d")
    if len(existing) < years * 4:
        print(f"数据不足（{len(existing)}条），拉取 {start_date} 至 {end_date} 的数据...")
        try:
            fetch_and_save(ts_code=ts_code, start_date=start_date, end_date=end_date)
            return True
        except Exception as e:
            print(f"拉取失败: {e}")
            return False
    return True

import os
import subprocess
import sys
from pathlib import Path

from config import DB_PATH


def parse_dates(value: str):
    if not value:
        return []
    parts = [p.strip() for p in value.split(",")]
    return [p for p in parts if p]


def run_cmd(cmd):
    subprocess.run(cmd, check=True)


def reset_db(db_path: Path):
    if db_path.exists():
        db_path.unlink()


def build_ingest_cmd(ingest_path: Path, date_str: str, args):
    cmd = [sys.executable, str(ingest_path), "--date", date_str]
    if args.report_rc_file:
        cmd += ["--report-rc-file", args.report_rc_file]
    if args.forecast_file:
        cmd += ["--forecast-file", args.forecast_file]
    if args.express_file:
        cmd += ["--express-file", args.express_file]
    if args.income_file:
        cmd += ["--income-file", args.income_file]
    if args.disclosure_file:
        cmd += ["--disclosure-file", args.disclosure_file]
    if args.only_report_rc:
        cmd.append("--only-report-rc")
    if args.no_vip:
        cmd.append("--no-vip")
    if args.env_path:
        cmd += ["--env-path", args.env_path]
    return cmd


def build_simple_cmd(script_path: Path):
    return [sys.executable, str(script_path)]


def resolve_db_path(root_dir: Path = None):
    """返回数据库路径（使用集中配置）"""
    return DB_PATH


def normalize_path(value: str):
    if not value:
        return value
    return os.path.abspath(value)

#!/usr/bin/env python3
"""
对比实际业绩与市场预期，生成告警
"""

import argparse
import json
from datetime import datetime

import pandas as pd

from db import connect, run_migrations, upsert_rows


def now_ts():
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")


def safe_ratio(actual, expected):
    if expected is None:
        return None
    if expected == 0:
        return None
    return (actual - expected) / abs(expected)


def load_table(conn, sql):
    return pd.read_sql_query(sql, conn)


def load_table_with_columns(conn, table, columns):
    info = pd.read_sql_query(f"PRAGMA table_info({table})", conn)
    available = set(info["name"].tolist())
    selected = [col for col in columns if col in available]
    if not selected:
        return pd.DataFrame(columns=columns)
    sql = f"SELECT {', '.join(selected)} FROM {table}"
    df = pd.read_sql_query(sql, conn)
    for col in columns:
        if col not in df.columns:
            df[col] = None
    return df


def build_expectations(conn, actuals):
    df = load_table(
        conn,
        """
        SELECT ts_code, period, report_date, np
        FROM report_rc
        WHERE np IS NOT NULL
        """
    )
    if df.empty or actuals.empty:
        return pd.DataFrame()

    df["period"] = df["period"].fillna(df["report_date"])
    df = df.dropna(subset=["period", "report_date"])
    df["period"] = df["period"].astype(str).apply(normalize_quarter_label)
    df["report_date"] = df["report_date"].astype(str)

    actuals = actuals.copy()
    actuals["period"] = actuals["period"].astype(str).apply(normalize_quarter_label)
    actuals["ann_date"] = actuals["ann_date"].astype(str)

    merged = actuals.merge(
        df,
        on=["ts_code", "period"],
        how="left"
    )
    merged["report_date_dt"] = pd.to_datetime(merged["report_date"], errors="coerce")
    merged["ann_date_dt"] = pd.to_datetime(merged["ann_date"], errors="coerce")
    merged = merged.dropna(subset=["report_date_dt", "ann_date_dt"])
    merged = merged[merged["report_date_dt"] < merged["ann_date_dt"]]
    merged = merged[merged["report_date_dt"] >= (merged["ann_date_dt"] - pd.Timedelta(days=365))]
    if merged.empty:
        return pd.DataFrame()

    agg = merged.groupby(["ts_code", "period", "ann_date"]).agg(
        expected_min=("np", "min"),
        expected_mean=("np", "mean"),
        expected_median=("np", "median"),
        expected_max=("np", "max"),
        expected_report_date_min=("report_date", "min"),
        expected_report_date_max=("report_date", "max"),
    ).reset_index()
    return agg


def latest_by_group(df, group_cols, order_col):
    if df.empty:
        return df
    df = df.sort_values(order_col, ascending=False)
    return df.groupby(group_cols, as_index=False).first()


def annualize_factor(end_date: str):
    if not end_date:
        return 1.0
    end_date = str(end_date)
    suffix = end_date[-4:]
    if suffix == "0331":
        return 4.0
    if suffix == "0630":
        return 2.0
    if suffix == "0930":
        return 4.0 / 3.0
    return 1.0


def end_date_to_annual_period(end_date: str):
    if not end_date:
        return None
    end_date = str(end_date)
    if len(end_date) < 4:
        return None
    year = end_date[:4]
    if not year.isdigit():
        return None
    return f"{year}Q4"


def normalize_quarter_label(value: str):
    if value is None:
        return value
    value = str(value)
    if "Q" in value:
        year = value[:4]
        quarter = value[-1]
        if year.isdigit() and quarter in {"1", "2", "3", "4"}:
            return f"{year}Q{quarter}"
    if len(value) == 8 and value.isdigit():
        year = value[:4]
        suffix = value[-4:]
        mapping = {"0331": "1", "0630": "2", "0930": "3", "1231": "4"}
        if suffix in mapping:
            return f"{year}Q{mapping[suffix]}"
    return value


def build_actuals(conn):
    # forecast: use avg(net_profit_min, net_profit_max)
    forecast = load_table_with_columns(
        conn,
        "forecast",
        ["ts_code", "end_date", "ann_date", "net_profit_min", "net_profit_max", "change_reason", "updated_at"]
    )
    if not forecast.empty:
        forecast["actual_value"] = (
            forecast["net_profit_min"].fillna(0) + forecast["net_profit_max"].fillna(0)
        ) / 2
        forecast = latest_by_group(forecast, ["ts_code", "end_date"], "updated_at")
        forecast["source"] = "forecast"
        forecast["reason_hint"] = forecast["change_reason"]

    # express: n_income（原单位：元，转换为万元）
    express = load_table_with_columns(
        conn,
        "express",
        ["ts_code", "end_date", "ann_date", "n_income", "perf_summary", "updated_at"]
    )
    if not express.empty:
        express = latest_by_group(express, ["ts_code", "end_date"], "updated_at")
        express["actual_value"] = express["n_income"] / 10000  # 元 → 万元
        express["source"] = "express"
        express["reason_hint"] = express["perf_summary"]

    # income: n_income（原单位：元，转换为万元）
    income = load_table_with_columns(
        conn,
        "income",
        ["ts_code", "end_date", "ann_date", "f_ann_date", "n_income", "updated_at"]
    )
    if not income.empty:
        income = latest_by_group(income, ["ts_code", "end_date"], "updated_at")
        income["actual_value"] = income["n_income"] / 10000  # 元 → 万元
        income["source"] = "income"
        income["ann_date"] = income["f_ann_date"].fillna(income["ann_date"])
        income["reason_hint"] = None

    # priority: forecast > express > income
    frames = []
    if not forecast.empty:
        frames.append(forecast[["ts_code", "end_date", "ann_date", "actual_value", "source", "reason_hint"]])
    if not express.empty:
        frames.append(express[["ts_code", "end_date", "ann_date", "actual_value", "source", "reason_hint"]])
    if not income.empty:
        frames.append(income[["ts_code", "end_date", "ann_date", "actual_value", "source", "reason_hint"]])

    if not frames:
        return pd.DataFrame()

    combined = pd.concat(frames, ignore_index=True)
    priority = {"forecast": 1, "express": 2, "income": 3}
    combined["priority"] = combined["source"].map(priority)
    combined = combined.sort_values("priority")
    combined = combined.groupby(["ts_code", "end_date"], as_index=False).first()
    combined["annualize_factor"] = combined["end_date"].apply(annualize_factor)
    combined["actual_annualized"] = combined["actual_value"] * combined["annualize_factor"]
    combined["period"] = combined["end_date"].apply(end_date_to_annual_period)
    return combined.drop(columns=["priority"])


def build_alerts(expectations, actuals):
    if expectations.empty or actuals.empty:
        return pd.DataFrame()

    merged = actuals.merge(expectations, on=["ts_code", "period", "ann_date"], how="left")
    merged = merged.dropna(subset=["expected_max", "actual_annualized"])

    # 超预期：actual > expected_max
    above = merged[merged["actual_annualized"] > merged["expected_max"]].copy()
    if not above.empty:
        above["alert_type"] = "above"

    # 低于预期：actual < expected_mean
    below = merged[merged["actual_annualized"] < merged["expected_mean"]].copy()
    if not below.empty:
        below["alert_type"] = "below"

    # 符合预期：expected_mean <= actual <= expected_max
    inline = merged[
        (merged["actual_annualized"] >= merged["expected_mean"]) &
        (merged["actual_annualized"] <= merged["expected_max"])
    ].copy()
    if not inline.empty:
        inline["alert_type"] = "inline"

    frames = [df for df in [above, below, inline] if not df.empty]
    alerts = pd.concat(frames, ignore_index=True) if frames else pd.DataFrame()
    if alerts.empty:
        return alerts

    alerts["delta_mean"] = alerts.apply(
        lambda r: safe_ratio(r["actual_annualized"], r["expected_mean"]), axis=1
    )
    alerts["delta_median"] = alerts.apply(
        lambda r: safe_ratio(r["actual_annualized"], r["expected_median"]), axis=1
    )
    alerts["delta_max"] = alerts.apply(
        lambda r: safe_ratio(r["actual_annualized"], r["expected_max"]), axis=1
    )
    return alerts


def main():
    parser = argparse.ArgumentParser(description="对比并生成告警")
    parser.add_argument("--dry-run", action="store_true", help="只统计不入库")
    args = parser.parse_args()

    conn = connect()
    run_migrations(conn)

    actuals = build_actuals(conn)
    expectations = build_expectations(conn, actuals)
    alerts = build_alerts(expectations, actuals)

    if not alerts.empty and not args.dry_run:
        rows = []
        ts_now = now_ts()
        for _, row in alerts.iterrows():
            payload = {
                "alert_type": row.get("alert_type"),  # above/below/inline
                "expected_min": row.get("expected_min"),
                "expected_mean": row.get("expected_mean"),
                "expected_median": row.get("expected_median"),
                "expected_max": row.get("expected_max"),
                "expected_report_date_min": row.get("expected_report_date_min"),
                "expected_report_date_max": row.get("expected_report_date_max"),
                "ann_date": row.get("ann_date"),
                "period": row.get("period"),
                "delta_mean": row.get("delta_mean"),
                "delta_median": row.get("delta_median"),
                "delta_max": row.get("delta_max"),
                "actual_original": row.get("actual_value"),
                "annualize_factor": row.get("annualize_factor"),
                "reason_hint": row.get("reason_hint"),
            }
            rows.append(
                {
                    "ts_code": row["ts_code"],
                    "end_date": row["end_date"],
                    "actual_value": row["actual_annualized"],
                    "expected_mean": row.get("expected_mean"),
                    "expected_median": row.get("expected_median"),
                    "expected_max": row.get("expected_max"),
                    "delta_mean": row.get("delta_mean"),
                    "delta_median": row.get("delta_median"),
                    "delta_max": row.get("delta_max"),
                    "source": row["source"],
                    "created_at": ts_now,
                    "payload_json": json.dumps(payload, ensure_ascii=True, sort_keys=True),
                }
            )
        upsert_rows(conn, "alerts", rows, ["ts_code", "end_date", "source", "actual_value"], preserve_on_update=["created_at"])

    conn.close()

    print(f"expectations: {len(expectations)}")
    print(f"actuals: {len(actuals)}")
    print(f"alerts: {len(alerts)}")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
SQLite 工具函数
"""

import argparse
import sqlite3
from pathlib import Path

from config import DB_PATH

SCHEMA_PATH = Path(__file__).with_name("schema.sql")


def connect(db_path: Path = DB_PATH) -> sqlite3.Connection:
    """连接数据库"""
    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    return conn


def run_migrations(conn: sqlite3.Connection):
    """执行建表脚本"""
    schema = SCHEMA_PATH.read_text(encoding="utf-8")
    conn.executescript(schema)
    conn.commit()


def fetch_existing(conn: sqlite3.Connection, table: str, unique_cols: list, row: dict):
    """查询已存在记录"""
    where = " AND ".join([f"{col} = ?" for col in unique_cols])
    sql = f"SELECT * FROM {table} WHERE {where} LIMIT 1"
    values = [row.get(col) for col in unique_cols]
    cur = conn.execute(sql, values)
    result = cur.fetchone()
    return dict(result) if result else None


def upsert_rows(conn: sqlite3.Connection, table: str, rows: list, unique_cols: list, preserve_on_update: list = None) -> int:
    """批量 upsert
    
    Args:
        preserve_on_update: 这些字段在更新时保留原值（仅 INSERT 时写入）
    """
    if not rows:
        return 0

    preserve_on_update = preserve_on_update or []

    all_cols = []
    for row in rows:
        for key in row.keys():
            if key not in all_cols:
                all_cols.append(key)

    placeholders = ", ".join(["?"] * len(all_cols))
    columns_sql = ", ".join(all_cols)
    conflict_sql = ", ".join(unique_cols)
    update_cols = [c for c in all_cols if c not in unique_cols]
    
    # preserve_on_update 的字段保留原值，其他字段用新值覆盖
    update_parts = []
    for c in update_cols:
        if c in preserve_on_update:
            update_parts.append(f"{c} = COALESCE({table}.{c}, excluded.{c})")
        else:
            update_parts.append(f"{c} = excluded.{c}")
    update_sql = ", ".join(update_parts)

    sql = (
        f"INSERT INTO {table} ({columns_sql}) "
        f"VALUES ({placeholders}) "
        f"ON CONFLICT({conflict_sql}) DO UPDATE SET {update_sql}"
    )

    values = []
    for row in rows:
        values.append([row.get(col) for col in all_cols])

    conn.executemany(sql, values)
    conn.commit()
    return len(rows)


def init_db():
    """初始化数据库"""
    conn = connect()
    run_migrations(conn)
    conn.close()


def main():
    parser = argparse.ArgumentParser(description="初始化数据库")
    parser.add_argument("--init", action="store_true", help="创建数据库与表")
    args = parser.parse_args()

    if args.init:
        init_db()
        print(f"数据库已初始化: {DB_PATH}")
    else:
        parser.print_help()


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
数据拉取与入库
"""

import argparse
import json
import os
from datetime import datetime
from pathlib import Path

import pandas as pd
import tushare as ts
from dotenv import load_dotenv

from db import connect, run_migrations, fetch_existing, upsert_rows


def load_env(env_path: str = None):
    """加载 .env"""
    if env_path:
        load_dotenv(env_path)
        return

    for parent in Path(__file__).resolve().parents:
        candidate = parent / ".env"
        if candidate.exists():
            load_dotenv(candidate)
            return


def init_tushare(env_path: str = None):
    """初始化 Tushare"""
    load_env(env_path)
    api_key = os.getenv("TUSHARE_API_KEY")
    if not api_key:
        raise RuntimeError("未找到 TUSHARE_API_KEY")
    return ts.pro_api(api_key)


def now_ts():
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")


def get_periods_for_date(date_str: str) -> list:
    """根据日期自动计算需要拉取的 period 列表
    
    规则：
    - 1月1日 - 4月30日：上年Q4(1231) + 当年Q1(0331)
    - 5月1日 - 8月31日：当年Q2(0630)
    - 9月1日 - 10月31日：当年Q3(0930)
    - 11月1日 - 12月31日：当年Q4(1231) 预披露
    """
    dt = datetime.strptime(date_str, "%Y%m%d")
    year = dt.year
    month = dt.month
    
    periods = []
    
    if 1 <= month <= 4:
        # 年报 + Q1 披露季
        periods.append(f"{year - 1}1231")  # 上年年报
        periods.append(f"{year}0331")       # 当年Q1
    elif 5 <= month <= 8:
        # 中报披露季
        periods.append(f"{year}0630")
    elif 9 <= month <= 10:
        # Q3 披露季
        periods.append(f"{year}0930")
    else:  # 11-12月
        # 年报预披露期
        periods.append(f"{year}1231")
    
    return periods


def safe_value(value):
    if pd.isna(value):
        return None
    return value


def to_payload(row: dict):
    return json.dumps(row, ensure_ascii=True, sort_keys=True)


def build_change_log(existing: dict, new_row: dict, fields: list):
    if not existing:
        return "new", None

    changes = {}
    for field in fields:
        old = existing.get(field)
        new = new_row.get(field)
        if old != new:
            changes[field] = {"old": old, "new": new}

    if not changes:
        return "same", None

    return "update", json.dumps(changes, ensure_ascii=True, sort_keys=True)


def normalize_rows(df: pd.DataFrame, base_map: dict, source: str, existing_fetcher=None, change_fields=None):
    rows = []
    ts_now = now_ts()

    for _, row in df.iterrows():
        row_dict = {k: safe_value(v) for k, v in row.to_dict().items()}
        new_row = {col: safe_value(row_dict.get(src)) for col, src in base_map.items()}
        new_row["source"] = source
        new_row["payload_json"] = to_payload(row_dict)
        new_row["updated_at"] = ts_now
        new_row["created_at"] = ts_now

        if existing_fetcher and change_fields:
            existing = existing_fetcher(new_row)
            change_type, change_log = build_change_log(existing, new_row, change_fields)
            new_row["change_type"] = change_type
            new_row["change_log"] = change_log
        rows.append(new_row)

    return rows


def fetch_report_rc_full(pro, report_date: str):
    return pro.report_rc(report_date=report_date)


def fetch_forecast_by_periods(pro, periods: list, use_vip: bool):
    """按 period 列表拉取业绩预告"""
    dfs = []
    for period in periods:
        if use_vip:
            df = pro.forecast_vip(period=period)
        else:
            df = pro.forecast(period=period)
        if df is not None and not df.empty:
            print(f"  forecast period={period}: {len(df)} 条")
            dfs.append(df)
    if not dfs:
        return pd.DataFrame()
    return pd.concat(dfs, ignore_index=True).drop_duplicates()


def fetch_express_by_periods(pro, periods: list, use_vip: bool):
    """按 period 列表拉取业绩快报"""
    dfs = []
    for period in periods:
        if use_vip:
            df = pro.express_vip(period=period)
        else:
            df = pro.express(period=period)
        if df is not None and not df.empty:
            print(f"  express period={period}: {len(df)} 条")
            dfs.append(df)
    if not dfs:
        return pd.DataFrame()
    return pd.concat(dfs, ignore_index=True).drop_duplicates()


def fetch_income_by_periods(pro, periods: list, use_vip: bool):
    """按 period 列表拉取正式业绩"""
    dfs = []
    for period in periods:
        if use_vip:
            df = pro.income_vip(period=period)
        else:
            df = pro.income(period=period)
        if df is not None and not df.empty:
            print(f"  income period={period}: {len(df)} 条")
            dfs.append(df)
    if not dfs:
        return pd.DataFrame()
    return pd.concat(dfs, ignore_index=True).drop_duplicates()


def fetch_disclosure_full(pro, date_str: str):
    dfs = []
    for key in ["ann_date", "pre_date", "actual_date"]:
        df = pro.disclosure_date(**{key: date_str})
        if df is not None and not df.empty:
            dfs.append(df)
    if not dfs:
        return pd.DataFrame()
    return pd.concat(dfs, ignore_index=True).drop_duplicates()


def main():
    parser = argparse.ArgumentParser(description="数据拉取与入库")
    parser.add_argument("--date", help="日期 (YYYYMMDD)，默认今天")
    parser.add_argument("--dry-run", action="store_true", help="只统计不入库")
    parser.add_argument("--no-vip", action="store_true", help="不使用 VIP 接口")
    parser.add_argument("--env-path", help="指定 .env 路径")
    parser.add_argument("--report-rc-file", help="使用本地 report_rc CSV 文件进行模拟")
    parser.add_argument("--forecast-file", help="使用本地 forecast CSV 文件进行模拟")
    parser.add_argument("--express-file", help="使用本地 express CSV 文件进行模拟")
    parser.add_argument("--income-file", help="使用本地 income CSV 文件进行模拟")
    parser.add_argument("--disclosure-file", help="使用本地 disclosure_date CSV 文件进行模拟")
    parser.add_argument("--only-report-rc", action="store_true", help="只处理 report_rc")
    args = parser.parse_args()

    date_str = args.date or datetime.now().strftime("%Y%m%d")
    use_vip = not args.no_vip
    
    # 根据日期自动计算需要拉取的 period
    periods = get_periods_for_date(date_str)
    print(f"日期: {date_str}")
    print(f"自动 period: {periods}")

    need_api = True
    if args.only_report_rc and args.report_rc_file:
        need_api = False
    if args.report_rc_file and args.forecast_file and args.express_file and args.income_file and args.disclosure_file:
        need_api = False

    pro = None if not need_api else init_tushare(args.env_path)
    conn = connect()
    run_migrations(conn)

    def existing_fetcher_factory(table, unique_cols):
        def _fetch(row):
            return fetch_existing(conn, table, unique_cols, row)
        return _fetch

    counts = {}

    # report_rc
    if args.report_rc_file:
        df_report = pd.read_csv(args.report_rc_file)
        if "report_date" in df_report.columns:
            df_report = df_report[df_report["report_date"].astype(str) == date_str]
    else:
        df_report = fetch_report_rc_full(pro, date_str)
    report_base = {
        "ts_code": "ts_code",
        "report_date": "report_date",
        "org_name": "org_name",
        "period": "quarter",
        "np": "np",
        "eps": "eps",
        "quarter": "quarter",
    }
    report_rows = normalize_rows(
        df_report,
        report_base,
        "report_rc",
        existing_fetcher_factory("report_rc", ["ts_code", "report_date", "org_name", "period"]),
        ["np"]
    )
    counts["report_rc"] = len(report_rows)
    if not args.dry_run and report_rows:
        upsert_rows(conn, "report_rc", report_rows, ["ts_code", "report_date", "org_name", "period"])

    if args.only_report_rc:
        df_forecast = pd.DataFrame()
    elif args.forecast_file:
        df_forecast = pd.read_csv(args.forecast_file)
        if "ann_date" in df_forecast.columns:
            df_forecast = df_forecast[df_forecast["ann_date"].astype(str) == date_str]
    else:
        print("拉取 forecast...")
        df_forecast = fetch_forecast_by_periods(pro, periods, use_vip)
    forecast_base = {
        "ts_code": "ts_code",
        "ann_date": "ann_date",
        "end_date": "end_date",
        "net_profit_min": "net_profit_min",
        "net_profit_max": "net_profit_max",
        "type": "type",
    }
    forecast_rows = normalize_rows(
        df_forecast,
        forecast_base,
        "forecast",
        existing_fetcher_factory("forecast", ["ts_code", "ann_date", "end_date"]),
        ["net_profit_min", "net_profit_max"]
    )
    counts["forecast"] = len(forecast_rows)
    if not args.dry_run and forecast_rows:
        upsert_rows(conn, "forecast", forecast_rows, ["ts_code", "ann_date", "end_date"])

    if args.only_report_rc:
        df_express = pd.DataFrame()
    elif args.express_file:
        df_express = pd.read_csv(args.express_file)
        if "ann_date" in df_express.columns:
            df_express = df_express[df_express["ann_date"].astype(str) == date_str]
    else:
        print("拉取 express...")
        df_express = fetch_express_by_periods(pro, periods, use_vip)
    express_base = {
        "ts_code": "ts_code",
        "ann_date": "ann_date",
        "end_date": "end_date",
        "n_income": "n_income",
    }
    express_rows = normalize_rows(
        df_express,
        express_base,
        "express",
        existing_fetcher_factory("express", ["ts_code", "ann_date", "end_date"]),
        ["n_income"]
    )
    counts["express"] = len(express_rows)
    if not args.dry_run and express_rows:
        upsert_rows(conn, "express", express_rows, ["ts_code", "ann_date", "end_date"])

    if args.only_report_rc:
        df_income = pd.DataFrame()
    elif args.income_file:
        df_income = pd.read_csv(args.income_file)
        if "ann_date" in df_income.columns:
            df_income = df_income[df_income["ann_date"].astype(str) == date_str]
    else:
        print("拉取 income...")
        df_income = fetch_income_by_periods(pro, periods, use_vip)
    income_base = {
        "ts_code": "ts_code",
        "ann_date": "ann_date",
        "end_date": "end_date",
        "n_income": "n_income",
        "report_type": "report_type",
        "comp_type": "comp_type",
    }
    income_rows = normalize_rows(
        df_income,
        income_base,
        "income",
        existing_fetcher_factory("income", ["ts_code", "ann_date", "end_date"]),
        ["n_income"]
    )
    counts["income"] = len(income_rows)
    if not args.dry_run and income_rows:
        upsert_rows(conn, "income", income_rows, ["ts_code", "ann_date", "end_date"])

    if args.only_report_rc:
        df_disclosure = pd.DataFrame()
    elif args.disclosure_file:
        df_disclosure = pd.read_csv(args.disclosure_file)
        date_cols = [c for c in ["ann_date", "pre_date", "actual_date"] if c in df_disclosure.columns]
        if date_cols:
            mask = False
            for col in date_cols:
                mask = mask | (df_disclosure[col].astype(str) == date_str)
            df_disclosure = df_disclosure[mask]
    else:
        df_disclosure = fetch_disclosure_full(pro, date_str)
    disclosure_base = {
        "ts_code": "ts_code",
        "end_date": "end_date",
        "pre_date": "pre_date",
        "ann_date": "ann_date",
        "actual_date": "actual_date",
    }
    disclosure_rows = normalize_rows(
        df_disclosure,
        disclosure_base,
        "disclosure_date",
        existing_fetcher_factory("disclosure_date", ["ts_code", "end_date"]),
        ["pre_date", "ann_date", "actual_date"]
    )
    counts["disclosure_date"] = len(disclosure_rows)
    if not args.dry_run and disclosure_rows:
        upsert_rows(conn, "disclosure_date", disclosure_rows, ["ts_code", "end_date"])

    conn.close()

    print(f"日期: {date_str}")
    print(f"使用VIP: {use_vip}")
    for key, value in counts.items():
        print(f"{key}: {value}")


if __name__ == "__main__":
    main()