#!/usr/bin/env python3
"""
backfill.py — 快讯历史数据补采工具
支持按时间范围补采缺失的历史数据

用法:
    python3 backfill.py --hours 24          # 补采过去24小时数据
    python3 backfill.py --date 2026-07-13   # 补采指定日期数据
    python3 backfill.py --hours 24 --source cls wallstreetcn  # 指定数据源
"""
import asyncio, json, os, sys, time, argparse, re
from datetime import datetime, timedelta, timezone
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from common import setup_kuaixun_paths, load_env, KX_HOME, KXS_HOME, BUNDLED_KXS_HOME, BIN_DIR, MYSQL_CFG, database_url
setup_kuaixun_paths()

_env_path = f"{BUNDLED_KXS_HOME or KXS_HOME}/.env"
load_env(_env_path)

import pymysql
from db_pool import DBContext, get_conn, put_conn

from extend.fetchers.manager import FetcherManager
from extend.fetchers import NewsItem

SOURCE_TABLE = {
    "sina": "news_sina",
    "cls": "news_cls",
    "wallstreetcn": "news_wallstreetcn",
    "jin10": "news_jin10",
    "eastmoney": "news_eastmoney",
    "36kr": "news_36kr",
    "gnews": "news_gnews",
    "gelonghui": "news_gelonghui",
    "newsapi": "news_newsapi",
    "hotlist": "news_hotlist",
    "21jingji": "news_21jingji",
    "yicai": "news_yicai",
    "futunn": "news_futunn",
    "hexun": "news_hexun",
    "investing": "news_investing",
    "jrj": "news_jrj",
    "jiemian": "news_jiemian",
    "xuangubao": "news_xuangubao",
    "xueqiu": "news_xueqiu",
    "10jqka": "news_10jqka",
}


def get_source_map() -> dict:
    # 连接必须在 finally 归还：原实现只在成功路径 put_conn，
    # 一旦 execute 抛错就会永久占住一个池位（池有上限，反复失败会耗尽）。
    conn = get_conn()
    try:
        cur = conn.cursor()
        cur.execute("SELECT id, code FROM news_sources WHERE is_active = 1")
        result = {r[1]: r[0] for r in cur.fetchall()}
        cur.close()
        return result
    finally:
        put_conn(conn)


def get_existing_ids(source_code: str, start_ts: int, end_ts: int) -> set:
    existing = set()
    conn = None
    try:
        conn = get_conn()
        cur = conn.cursor()
        per_table = SOURCE_TABLE.get(source_code)
        if per_table:
            cur.execute(
                f"SELECT source_news_id FROM {per_table} WHERE published_at >= %s AND published_at <= %s",
                (start_ts, end_ts),
            )
            for row in cur.fetchall():
                existing.add(row[0])
        cur.close()
    except Exception as e:
        print(f"  [backfill] 获取已存在ID失败: {e}", flush=True)
    finally:
        # 同上：原实现把 put_conn 放在 try 内，异常时连接不会归还
        if conn is not None:
            put_conn(conn)
    return existing


def insert_news(items: list[tuple]) -> dict:
    conn = get_conn()
    cur = conn.cursor()
    stats = {"total": 0, "new": 0, "dup": 0}
    now_ts = int(time.time() * 1000)
    try:
        for row in items:
          sid, scode, nid, title, content, summary, url, cat, tags, imp, pub = row
          title_s = (title or "")[:500]
          content_s = (content or "")[:2000]
          summary_s = (summary or "")[:1000]
          url_s = (url or "")[:500]

          try:
              cur.execute(
                  """INSERT IGNORE INTO news_flash
                  (source_id, source_news_id, title, content, summary, url,
                   category, tags, importance, published_at, fetched_at, is_duplicate)
                  VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, NOW(), 0)""",
                  (sid, nid, title_s, content_s, summary_s, url_s,
                   cat or "general", json.dumps(tags, ensure_ascii=False) if tags else "[]",
                   imp or 0, pub or now_ts),
              )
          except Exception as e:
              print(f"  ⚠️ 写入 news_flash 失败 [{scode}:{nid}]: {e}", flush=True)

          per_table = SOURCE_TABLE.get(scode)
          if per_table:
              try:
                  cur.execute(
                      f"""INSERT IGNORE INTO {per_table}
                      (source_news_id, title, content, url, published_at,
                       fetched_at, importance, is_dup, summary, tags, category)
                      VALUES (%s, %s, %s, %s, %s, NOW(), %s, 0, %s, %s, %s)""",
                      (nid, title_s, content_s, url_s, pub or now_ts, imp or 0, summary_s,
                       json.dumps(tags, ensure_ascii=False) if tags else "[]",
                       cat or "general"),
                  )
                  if cur.rowcount > 0:
                      stats["new"] += 1
                  else:
                      stats["dup"] += 1
                  stats["total"] += 1
              except Exception as e:
                  print(f"  ⚠️ 写入 {per_table} 失败 [{nid}]: {e}", flush=True)

        conn.commit()
        return stats
    finally:
        # 嵌套 try/finally 是必需的：若 cur.close() 本身抛错（连接已断时很常见），
        # 直接写成 `cur.close(); put_conn(conn)` 会跳过归还 → 池位照样泄漏。
        # 归还连接必须是无条件的最后一步。
        try:
            cur.close()
        finally:
            put_conn(conn)


class PagedFetcher:
    @staticmethod
    async def fetch_cls_paged(page: int, headers: dict) -> list:
        import httpx
        url = f"https://newsnow.busiyi.world/api/s?id=cls&latest&page={page}"
        items = []
        try:
            async with httpx.AsyncClient(timeout=15) as client:
                resp = await client.get(url, headers=headers)
                if resp.status_code == 200:
                    data = resp.json()
                    for item in data.get("items", []):
                        news_id = str(item.get("id", ""))
                        title = item.get("title", "")
                        items.append({
                            "source_news_id": news_id,
                            "title": title,
                            "content": title,
                            "summary": title,
                            "url": item.get("url", f"https://www.cls.cn/detail/{news_id}"),
                            "published_at": item.get("pubDate", 0),
                        })
        except Exception as e:
            print(f"  [cls] 分页{page}失败: {e}", flush=True)
        return items

    @staticmethod
    async def fetch_eastmoney_paged(page: int, headers: dict) -> list:
        import httpx
        url = f"https://newsapi.eastmoney.com/kuaixun/v1/getlist_102_ajaxResult_50_{page}_.html"
        items = []
        try:
            async with httpx.AsyncClient(timeout=15) as client:
                resp = await client.get(url, headers=headers)
                match = re.search(r"var ajaxResult\s*=\s*(.+?)\s*$", resp.text, re.DOTALL)
                if match:
                    data = json.loads(match.group(1))
                    for item in data.get("LivesList", []):
                        news_id = str(item.get("newsid", ""))
                        title = item.get("title", "").strip()
                        content = item.get("digest", "").strip()
                        show_time = item.get("showtime", "")
                        try:
                            dt = datetime.strptime(show_time, "%Y-%m-%d %H:%M:%S")
                            pub_ts = int(dt.timestamp() * 1000)
                        except ValueError:
                            pub_ts = int(time.time() * 1000)
                        items.append({
                            "source_news_id": news_id,
                            "title": title,
                            "content": content,
                            "summary": content[:200] if content else "",
                            "url": item.get("url_w", ""),
                            "published_at": pub_ts,
                        })
        except Exception as e:
            print(f"  [eastmoney] 分页{page}失败: {e}", flush=True)
        return items

    @staticmethod
    async def fetch_10jqka_paged(page: int, headers: dict) -> list:
        import httpx
        url = "https://news.10jqka.com.cn/tapp/news/push/stock/"
        items = []
        try:
            async with httpx.AsyncClient(timeout=15) as client:
                resp = await client.get(url, params={"page": page, "tag": "", "track": "website"}, headers=headers)
                data = resp.json()
                for item in data.get("data", {}).get("list", []):
                    news_id = str(item.get("seq", "") or item.get("id", ""))
                    title = item.get("title", "").strip()
                    content = item.get("digest", "").strip() or item.get("short", "").strip()
                    ctime = int(item.get("ctime", 0))
                    pub_ts = ctime * 1000 if ctime else int(time.time() * 1000)
                    items.append({
                        "source_news_id": news_id,
                        "title": title,
                        "content": content,
                        "summary": content[:200] if content else "",
                        "url": item.get("url", ""),
                        "published_at": pub_ts,
                    })
        except Exception as e:
            print(f"  [10jqka] 分页{page}失败: {e}", flush=True)
        return items

    @staticmethod
    async def fetch_21jingji_paged(page: int, headers: dict) -> list:
        import httpx
        url = "https://api.21jingji.com/timestream/getListweb"
        items = []
        try:
            async with httpx.AsyncClient(timeout=15) as client:
                resp = await client.get(url, params={"page": page}, headers=headers)
                match = re.search(r"\{.*\}", resp.text, re.DOTALL)
                if match:
                    data = json.loads(match.group())
                    for item in data.get("list", []):
                        news_id = str(item.get("id", ""))
                        title = (item.get("title") or "").strip()
                        content = (item.get("content") or "").strip()
                        if not title:
                            title = content[:80]
                        pub_time = item.get("inputtime", "")
                        try:
                            dt = datetime.strptime(pub_time, "%Y-%m-%d %H:%M:%S")
                            pub_ts = int(dt.timestamp() * 1000)
                        except ValueError:
                            pub_ts = int(time.time() * 1000)
                        items.append({
                            "source_news_id": news_id,
                            "title": title,
                            "content": content,
                            "summary": content[:200],
                            "url": item.get("url", ""),
                            "published_at": pub_ts,
                        })
        except Exception as e:
            print(f"  [21jingji] 分页{page}失败: {e}", flush=True)
        return items

    @staticmethod
    async def fetch_yicai_paged(page: int, headers: dict) -> list:
        import httpx
        url = "https://www.yicai.com/api/ajax/getbrieflist"
        items = []
        try:
            async with httpx.AsyncClient(timeout=15) as client:
                resp = await client.get(url, params={"page": page, "pagesize": 20, "type": "0", "id": "0"}, headers=headers)
                for item in resp.json():
                    news_id = str(item.get("LiveID", "") or item.get("NewsID", ""))
                    title = (item.get("LiveTitle") or item.get("NewsTitle") or "").strip()
                    content = (item.get("LiveContent") or "").replace("<b>", "").replace("</b>", "").strip()
                    if not title:
                        title = content[:80]
                    pub_time = item.get("CreateDate", "")
                    try:
                        dt = datetime.strptime(pub_time, "%Y-%m-%d %H:%M:%S")
                        pub_ts = int(dt.timestamp() * 1000)
                    except ValueError:
                        pub_ts = int(time.time() * 1000)
                    url_val = item.get("ShareUrl", "")
                    if not url_val and item.get("url"):
                        url_val = f"https://www.yicai.com{item['url']}"
                    items.append({
                        "source_news_id": news_id,
                        "title": title,
                        "content": content,
                        "summary": content[:200],
                        "url": url_val,
                        "published_at": pub_ts,
                    })
        except Exception as e:
            print(f"  [yicai] 分页{page}失败: {e}", flush=True)
        return items

    @staticmethod
    async def fetch_hexun_paged(page: int, headers: dict) -> list:
        import httpx
        url = "https://opentool.hexun.com/MongodbNewsService/getNewsListByJson.jsp"
        callback = f"ptemplate_jsonp_{int(time.time() * 1000000)}"
        items = []
        try:
            async with httpx.AsyncClient(timeout=15) as client:
                resp = await client.get(url, params={"id": "189223574", "s": 20, "cp": page, "callback": callback}, headers=headers)
                start_idx = resp.text.find("(")
                end_idx = resp.text.rfind(")")
                if start_idx != -1 and end_idx != -1:
                    data = json.loads(resp.text[start_idx + 1:end_idx].strip())
                    for item in data.get("result", []):
                        news_id = str(item.get("id", ""))
                        title = (item.get("title") or "").strip()
                        content = (item.get("abstract") or "").strip()
                        if not title:
                            title = content[:80]
                        pub_time = item.get("entitytime", "")
                        try:
                            dt = datetime.strptime(pub_time, "%Y-%m-%d %H:%M:%S")
                            pub_ts = int(dt.timestamp() * 1000)
                        except ValueError:
                            pub_ts = int(time.time() * 1000)
                        items.append({
                            "source_news_id": news_id,
                            "title": title,
                            "content": content,
                            "summary": content[:200],
                            "url": item.get("entityurl", ""),
                            "published_at": pub_ts,
                        })
        except Exception as e:
            print(f"  [hexun] 分页{page}失败: {e}", flush=True)
        return items


async def backfill_source_paged(source_code: str, start_ts: int, end_ts: int, source_map: dict) -> dict:
    """使用分页API补采单个数据源"""
    stats = {"source": source_code, "total": 0, "new": 0, "dup": 0, "pages": 0}
    existing_ids = get_existing_ids(source_code, start_ts, end_ts)
    print(f"  [{source_code}] 已存在 {len(existing_ids)} 条，开始分页补采...", flush=True)

    sid = source_map.get(source_code)
    if sid is None:
        print(f"  [{source_code}] 数据源未注册，跳过", flush=True)
        return stats

    all_items = []
    page = 1
    max_pages = 30
    seen_ids = set()

    headers = {
        "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
        "Referer": "https://www.findatahub.com/",
    }

    fetcher_func = None
    if source_code == "cls":
        headers["Referer"] = "https://www.cls.cn/"
        fetcher_func = PagedFetcher.fetch_cls_paged
    elif source_code == "eastmoney":
        headers["Referer"] = "https://kuaixun.eastmoney.com/"
        fetcher_func = PagedFetcher.fetch_eastmoney_paged
    elif source_code == "10jqka":
        headers["Referer"] = "https://news.10jqka.com.cn/"
        fetcher_func = PagedFetcher.fetch_10jqka_paged
    elif source_code == "21jingji":
        headers["Referer"] = "https://www.21jingji.com/"
        fetcher_func = PagedFetcher.fetch_21jingji_paged
    elif source_code == "yicai":
        headers["Referer"] = "https://www.yicai.com/brief/"
        fetcher_func = PagedFetcher.fetch_yicai_paged
    elif source_code == "hexun":
        headers["Referer"] = "https://www.hexun.com/"
        fetcher_func = PagedFetcher.fetch_hexun_paged

    if not fetcher_func:
        print(f"  [{source_code}] 不支持分页补采，跳过", flush=True)
        return stats

    while page <= max_pages:
        try:
            items = await fetcher_func(page, headers)
            if not items:
                break

            new_items = []
            for item in items:
                nid = item.get("source_news_id", "") or item.get("url", "")[-40:]
                if nid in seen_ids:
                    continue
                seen_ids.add(nid)

                if nid in existing_ids:
                    continue

                pub_ts = item.get("published_at", 0)
                if pub_ts < start_ts or pub_ts > end_ts:
                    continue

                new_items.append({
                    "sid": sid,
                    "source_code": source_code,
                    "source_news_id": nid,
                    "title": item.get("title", ""),
                    "content": item.get("content", ""),
                    "summary": item.get("summary", ""),
                    "url": item.get("url", ""),
                    "category": "general",
                    "tags": [],
                    "importance": 0,
                    "published_at": pub_ts,
                })

            all_items.extend(new_items)
            stats["pages"] += 1

            if len(new_items) == 0:
                break

            print(f"  [{source_code}] 第{page}页: 新增{len(new_items)}条（累计{len(all_items)}条）", flush=True)

            page += 1
            await asyncio.sleep(1.5)

        except Exception as e:
            print(f"  [{source_code}] 第{page}页失败: {e}", flush=True)
            break

    if all_items:
        rows = []
        for item in all_items:
            rows.append((
                item["sid"], item["source_code"], item["source_news_id"],
                item["title"], item["content"], item.get("summary") or "",
                item["url"], item["category"], item["tags"],
                item.get("importance") or 0, item["published_at"],
            ))
        r = insert_news(rows)
        stats["total"] = r["total"]
        stats["new"] = r["new"]
        stats["dup"] = r["dup"]
        print(f"  [{source_code}] 分页补采完成: 新增{r['new']}条, 重复{r['dup']}条", flush=True)
    else:
        print(f"  [{source_code}] 无新增数据", flush=True)

    return stats


async def backfill_all(start_ts: int, end_ts: int, sources: list = None) -> dict:
    """补采所有数据源"""
    source_map = get_source_map()

    paged_sources = ["cls", "eastmoney", "10jqka", "21jingji", "yicai", "hexun"]
    active_sources = sources if sources else paged_sources

    print(f"\n🚀 开始分页补采 {start_ts} - {end_ts} 的数据", flush=True)
    print(f"📊 数据源: {', '.join(active_sources)}", flush=True)

    all_stats = []
    for source_code in active_sources:
        if source_code not in paged_sources:
            print(f"\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", flush=True)
            print(f"  [{source_code}] 不支持分页补采，跳过", flush=True)
            continue

        print(f"\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", flush=True)
        r = await backfill_source_paged(source_code, start_ts, end_ts, source_map)
        all_stats.append(r)

    total_new = sum(s["new"] for s in all_stats)
    total_dup = sum(s["dup"] for s in all_stats)

    print(f"\n{'='*50}", flush=True)
    print(f"📝 分页补采完成!", flush=True)
    print(f"   新增: {total_new} 条", flush=True)
    print(f"   重复: {total_dup} 条", flush=True)
    print(f"{'='*50}", flush=True)

    return {"total_new": total_new, "total_dup": total_dup, "sources": all_stats}


def main():
    parser = argparse.ArgumentParser(description="快讯历史数据补采工具")
    parser.add_argument("--hours", type=int, help="补采过去N小时数据")
    parser.add_argument("--date", type=str, help="补采指定日期数据 (格式: YYYY-MM-DD)")
    parser.add_argument("--source", nargs="+", help="指定数据源 (如: cls eastmoney)")
    args = parser.parse_args()

    if args.hours:
        end_ts = int(time.time() * 1000)
        start_ts = end_ts - args.hours * 3600 * 1000
        print(f"📅 补采范围: 过去{args.hours}小时", flush=True)
    elif args.date:
        try:
            dt = datetime.strptime(args.date, "%Y-%m-%d")
            start_ts = int(dt.timestamp() * 1000)
            end_ts = start_ts + 24 * 3600 * 1000 - 1
            print(f"📅 补采范围: {args.date}", flush=True)
        except ValueError:
            print("❌ 日期格式错误，请使用 YYYY-MM-DD", flush=True)
            sys.exit(1)
    else:
        print("❌ 请指定 --hours 或 --date 参数", flush=True)
        parser.print_help()
        sys.exit(1)

    print(f"⌚ 时间范围: {start_ts} -> {end_ts}", flush=True)

    stats = asyncio.run(backfill_all(start_ts, end_ts, args.source))
    print(json.dumps(stats, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()