Jiucai Capture

# 韭菜公社平台登录账号（手机号）
JIUCAI_PHONE=your_phone_number_here

# 韭菜公社平台登录密码
JIUCAI_PASSWORD=your_password_here

__pycache__/
*.pyc
.venv/
*.egg-info/
dist/
.pytest_cache/
data/*.db
data/*.log
data/*.txt
data/screenshots/
.DS_Store
.env

"""A股主题跟踪 MCP Server"""

"""python -m a_stock_watcher 入口"""
from .server import main

main()

"""AI 解析层 — 调用 Gemini 2.5 Flash 提取股票、主题、投资逻辑，同时过滤噪音"""

import os
import yaml
from google import genai

from .models import ParsedResult, StockMention, ThemeMention

# Gemini 客户端（延迟初始化）
_client: genai.Client | None = None


def _get_client() -> genai.Client:
    global _client
    if _client is None:
        api_key = os.environ.get("GEMINI_API_KEY", "")
        if not api_key:
            raise ValueError("GEMINI_API_KEY 环境变量未设置")
        _client = genai.Client(api_key=api_key)
    return _client


PARSE_PROMPT = """你是一个 A 股投研分析师。请分析以下文章内容，提取结构化信息。

## 任务
1. **评估相关性** (relevance, 0-10)：与 A 股投资是否相关？广告/软文/水文/无实质分析 → 给低分。
2. **提取发布日期** (publish_date)
3. **提取涉及的股票** (stocks)：代码、名称、上下文片段、投资逻辑
4. **提取涉及的主题** (themes)：主题名称、分类（消费/科技/医药/周期/金融/新能源 等）
5. **总结投资逻辑** (logic_summary)

## 输出要求
请严格输出以下 YAML 格式，不要额外的 markdown 包裹：

relevance: <0-10整数>
publish_date: "<YYYY-MM-DD或空>"
stocks:
  - code: "<6位股票代码>"
    name: "<股票名称>"
    context: "<提及该股票的上下文片段（50字内）>"
    logic: "<该股票的投资逻辑（50字内）>"
themes:
  - name: "<主题名称>"
    category: "<分类>"
logic_summary: "<整体投资逻辑总结（100字内）>"

如果文章无关 A 股投资，relevance 设为低分，其余字段可以为空列表。

## 文章标题
{title}

## 文章内容
{content}
"""


async def parse_article(title: str, content: str) -> ParsedResult:
    """
    调用 Gemini 2.5 Flash 解析文章，返回结构化结果。
    
    - relevance < 5 → 标记为噪音，调用方决定是否入库
    - 解析失败 → parse_failed=True，保留原文
    """
    if not content.strip():
        return ParsedResult(
            relevance=0, parse_failed=True,
            filter_reason="空内容",
        )

    prompt = PARSE_PROMPT.format(title=title, content=content[:8000])

    try:
        client = _get_client()
        response = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=prompt,
        )
        raw_text = response.text.strip()

        # 清理可能的 markdown 包裹
        if raw_text.startswith("```"):
            lines = raw_text.split("\n")
            raw_text = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])

        parsed_yaml = yaml.safe_load(raw_text)
        if not isinstance(parsed_yaml, dict):
            return ParsedResult(
                relevance=0, raw_yaml=raw_text,
                parse_failed=True, filter_reason="AI 输出不是有效 YAML dict",
            )

        relevance = int(parsed_yaml.get("relevance", 0))
        result = ParsedResult(
            relevance=relevance,
            publish_date=str(parsed_yaml.get("publish_date", "")),
            logic_summary=str(parsed_yaml.get("logic_summary", "")),
            raw_yaml=raw_text,
            parse_failed=False,
        )

        # 解析股票
        for s in parsed_yaml.get("stocks", []) or []:
            if isinstance(s, dict) and s.get("code"):
                result.stocks.append(StockMention(
                    code=str(s["code"]),
                    name=str(s.get("name", "")),
                    context=str(s.get("context", "")),
                    logic=str(s.get("logic", "")),
                ))

        # 解析主题
        for t in parsed_yaml.get("themes", []) or []:
            if isinstance(t, dict) and t.get("name"):
                result.themes.append(ThemeMention(
                    name=str(t["name"]),
                    category=str(t.get("category", "")),
                ))

        # 噪音标记
        if relevance < 5:
            result.filter_reason = f"低相关性 relevance={relevance}"

        return result

    except Exception as e:
        return ParsedResult(
            relevance=0, parse_failed=True,
            filter_reason=f"AI 解析异常: {str(e)}",
        )


IMAGE_PARSE_PROMPT = """你是一个 A 股投研分析师。请分析这张产业链图片/图表，提取结构化信息。

## 任务
1. 识别图中涉及的**所有股票**（代码+名称）
2. 识别图中展示的**产业链结构**和**投资逻辑**
3. 提取关键数据点

## 输出要求
请严格输出以下 YAML 格式，不要额外的 markdown 包裹：

relevance: 8
publish_date: ""
stocks:
  - code: "<6位股票代码>"
    name: "<股票名称>"
    context: "<在图中的角色/位置>"
    logic: "<投资逻辑>"
themes:
  - name: "<产业/主题名称>"
    category: "<分类>"
logic_summary: "<整体产业链逻辑总结（150字内）>"

图片对应的产业名称: {industry_name}
"""


async def parse_image(image_path: str, industry_name: str = "") -> ParsedResult:
    """
    调用 Gemini Vision 分析图片，提取股票名单和产业链逻辑。

    Args:
        image_path: 截图文件的绝对路径
        industry_name: 产业链名称（上下文）
    """
    try:
        from google.genai import types
        client = _get_client()

        # 读取图片文件
        with open(image_path, "rb") as f:
            image_data = f.read()

        prompt = IMAGE_PARSE_PROMPT.format(industry_name=industry_name)

        response = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=[
                types.Part.from_bytes(data=image_data, mime_type="image/png"),
                prompt,
            ],
        )
        raw_text = response.text.strip()

        # 清理可能的 markdown 包裹
        if raw_text.startswith("```"):
            lines = raw_text.split("\n")
            raw_text = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])

        parsed_yaml = yaml.safe_load(raw_text)
        if not isinstance(parsed_yaml, dict):
            return ParsedResult(
                relevance=5, raw_yaml=raw_text,
                parse_failed=True, filter_reason="图片AI输出不是有效YAML",
            )

        relevance = int(parsed_yaml.get("relevance", 5))
        result = ParsedResult(
            relevance=relevance,
            publish_date=str(parsed_yaml.get("publish_date", "")),
            logic_summary=str(parsed_yaml.get("logic_summary", "")),
            raw_yaml=raw_text,
            parse_failed=False,
        )

        for s in parsed_yaml.get("stocks", []) or []:
            if isinstance(s, dict) and s.get("code"):
                result.stocks.append(StockMention(
                    code=str(s["code"]),
                    name=str(s.get("name", "")),
                    context=str(s.get("context", "")),
                    logic=str(s.get("logic", "")),
                ))

        for t in parsed_yaml.get("themes", []) or []:
            if isinstance(t, dict) and t.get("name"):
                result.themes.append(ThemeMention(
                    name=str(t["name"]),
                    category=str(t.get("category", "")),
                ))

        return result

    except Exception as e:
        return ParsedResult(
            relevance=0, parse_failed=True,
            filter_reason=f"图片AI解析异常: {str(e)}",
        )

"""登录工具 — 自动登录韭菜公社，保存登录态

用法（自动登录）：
    cd /path/to/capture-韭菜公社 && source .env && uv run python -m a_stock_watcher.auth

会打开浏览器 → 自动输入手机号和密码 → 登录 → 保存 cookie 到 data/auth_state.json。
后续所有爬虫运行会自动加载此文件，无需重复登录。

如需手动登录（自动登录失败时）：
    uv run python -m a_stock_watcher.auth --manual
"""

import os
import sys
import asyncio
import logging
from pathlib import Path
from playwright.async_api import async_playwright

logger = logging.getLogger("a_stock_watcher.auth")

AUTH_DIR = Path(__file__).parent.parent / "data"
AUTH_STATE_PATH = AUTH_DIR / "auth_state.json"


def _load_credentials() -> tuple[str, str]:
    """三级加载凭据：环境变量 > skill目录.env > 上级目录.env，全部缺失则报错退出"""
    phone = os.environ.get("JIUCAI_PHONE") or os.environ.get("JIUCAI_ACCOUNT")
    password = os.environ.get("JIUCAI_PASSWORD")
    if phone and password:
        return phone, password

    skill_dir = Path(__file__).parent.parent
    for env_path in [skill_dir / ".env", skill_dir.parent / ".env"]:
        if not env_path.exists():
            continue
        for line in env_path.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, _, v = line.partition("=")
            k = k.strip()
            v = v.strip().strip('"').strip("'")
            if k in ("JIUCAI_PHONE", "JIUCAI_ACCOUNT"):
                phone = v
            elif k == "JIUCAI_PASSWORD":
                password = v
        if phone and password:
            return phone, password

    print("❌ 未找到 JIUCAI_PHONE / JIUCAI_PASSWORD，请设置环境变量或创建 .env 文件", file=sys.stderr)
    sys.exit(1)


# 账号信息：优先环境变量，其次 .env 文件
PHONE, PASSWORD = _load_credentials()


async def auto_login():
    """自动登录：打开浏览器 → 输入账号密码 → 保存登录态"""
    AUTH_DIR.mkdir(parents=True, exist_ok=True)

    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=False)
        context = await browser.new_context()
        page = await context.new_page()

        print("🔄 正在打开韭菜公社...")
        await page.goto("https://www.jiuyangongshe.com", wait_until="networkidle", timeout=60000)
        await page.wait_for_timeout(2000)

        # 触发登录弹窗 — 点击需要登录的功能
        print("🔄 触发登录弹窗...")
        try:
            await page.click('text=交易计划', timeout=5000)
        except Exception:
            # 备选方案：直接跳转到需要登录的页面
            await page.goto("https://www.jiuyangongshe.com/tradingPlan", timeout=30000)
        await page.wait_for_timeout(3000)

        # 切换到"账号密码登录"
        pwd_tab = await page.query_selector('text=账号密码登录')
        if pwd_tab:
            await pwd_tab.click()
            await page.wait_for_timeout(1000)
            print("✅ 已切换到账号密码登录")
        else:
            print("⚠️  未找到'账号密码登录'，尝试继续...")

        # 查找并填写手机号
        phone_filled = False
        phone_selectors = [
            'input[placeholder*="手机"]',
            'input[placeholder*="账号"]',
            'input[type="tel"]',
        ]
        for sel in phone_selectors:
            phone_input = await page.query_selector(sel)
            if phone_input:
                await phone_input.click()
                await phone_input.fill("")
                await phone_input.fill(PHONE)
                phone_filled = True
                print(f"✅ 已输入手机号: {PHONE[:3]}****{PHONE[-4:]}")
                break

        if not phone_filled:
            # 回退：查找弹窗内所有非密码输入框
            inputs = await page.query_selector_all('.el-dialog input:not([type="password"]):not([type="hidden"])')
            if inputs:
                await inputs[0].click()
                await inputs[0].fill("")
                await inputs[0].fill(PHONE)
                phone_filled = True
                print(f"✅ 已输入手机号（回退选择器）")

        if not phone_filled:
            print("❌ 未找到手机号输入框！请使用 --manual 手动登录")
            await browser.close()
            return False

        # 查找并填写密码
        pwd_input = await page.query_selector('input[type="password"]')
        if pwd_input:
            await pwd_input.click()
            await pwd_input.fill("")
            await pwd_input.fill(PASSWORD)
            print("✅ 已输入密码")
        else:
            print("❌ 未找到密码输入框！请使用 --manual 手动登录")
            await browser.close()
            return False

        # 点击登录按钮
        login_clicked = False
        login_selectors = [
            '.el-dialog button:has-text("登录")',
            'button:has-text("登录")',
            '.login-btn',
        ]
        for sel in login_selectors:
            btn = await page.query_selector(sel)
            if btn and await btn.is_visible():
                await btn.click()
                login_clicked = True
                print("✅ 已点击登录")
                break

        if not login_clicked:
            # 回退：按回车
            await pwd_input.press("Enter")
            print("✅ 已按回车提交登录")

        # 等待登录完成
        await page.wait_for_timeout(5000)

        # 验证登录状态
        cookies = await context.cookies()
        cookie_names = [c['name'] for c in cookies]
        has_auth = any('token' in n.lower() or 'session' in n.lower()
                       or 'user' in n.lower() or 'auth' in n.lower()
                       for n in cookie_names)

        if has_auth or len(cookies) > 5:
            # 保存登录态
            await context.storage_state(path=str(AUTH_STATE_PATH))
            print(f"\n🎉 登录成功！登录态已保存到 {AUTH_STATE_PATH}")
            print(f"   Cookies 数量: {len(cookies)}")
            await browser.close()
            return True
        else:
            print(f"\n⚠️  登录可能未成功（Cookies: {len(cookies)}）")
            print("请检查浏览器窗口，手动完成登录后按回车保存...")
            input("\n按回车保存登录态...")
            await context.storage_state(path=str(AUTH_STATE_PATH))
            print(f"✅ 登录态已保存到 {AUTH_STATE_PATH}")
            await browser.close()
            return True


async def manual_login():
    """手动登录：打开浏览器，等待用户手动操作"""
    AUTH_DIR.mkdir(parents=True, exist_ok=True)

    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=False)
        context = await browser.new_context()
        page = await context.new_page()

        await page.goto("https://www.jiuyangongshe.com", wait_until="networkidle")
        print("\n" + "=" * 50)
        print("浏览器已打开，请手动完成登录。")
        print("登录成功后，回到这里按 回车 保存登录态。")
        print("=" * 50)

        input("\n按回车保存登录态...")

        await context.storage_state(path=str(AUTH_STATE_PATH))
        print(f"\n✅ 登录态已保存到 {AUTH_STATE_PATH}")
        await browser.close()


if __name__ == "__main__":
    if "--manual" in sys.argv:
        asyncio.run(manual_login())
    else:
        asyncio.run(auto_login())

"""CLI 入口 — 供 cron / launchd 等外部调度器调用

用法：
    # 抓取全部数据源
    uv run python -m a_stock_watcher.cli fetch

    # 抓取指定来源
    uv run python -m a_stock_watcher.cli fetch --source study_hot

    # 回补全文（对已有短内容文章用 crawl4ai 拉全文 + 重新 AI 解析）
    uv run python -m a_stock_watcher.cli backfill

    # 查看统计
    uv run python -m a_stock_watcher.cli stats

crontab 示例（每小时执行）：
    0 * * * * cd /path/to/capture-韭菜公社 && uv run python -m a_stock_watcher.cli fetch >> data/cron.log 2>&1
"""

import argparse
import asyncio
import json
import logging
import os
import sys
from datetime import datetime
from pathlib import Path


def _load_dotenv():
    """从 skill 目录或上级目录加载 .env（不依赖第三方库）"""
    skill_dir = Path(__file__).parent.parent
    for env_path in [skill_dir / ".env", skill_dir.parent / ".env"]:
        if not env_path.exists():
            continue
        for line in env_path.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, _, v = line.partition("=")
            k = k.strip()
            v = v.strip().strip('"').strip("'")
            if k and k not in os.environ:  # 已有环境变量的不覆盖
                os.environ[k] = v


_load_dotenv()

from .sources import SOURCES
from . import database
from .ai_parser import parse_article

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger("a_stock_watcher.cli")


async def cmd_fetch(source: str | None = None):
    """抓取→AI解析→入库"""
    targets = {source: SOURCES[source]} if source else SOURCES
    logger.info(f"开始抓取: {list(targets.keys())}")

    for name, scraper_cls in targets.items():
        try:
            scraper = scraper_cls()
            articles = await scraper.run()
            saved, skipped, filtered = 0, 0, 0

            updated = 0
            for article in articles:
                # 产业链：跳过前置去重，让 save_article 内部做 upsert（保留最新版本）
                # 其他来源：先查 DB，已存在的跳过，不浪费 AI 配额
                if article.source != "industry_chain":
                    if await database.check_exists(article.content_hash):
                        skipped += 1
                        continue

                # 如果已有解析结果，直接使用；否则调 Gemini 解析
                if article._parsed is not None:
                    parsed = article._parsed
                else:
                    parsed = await parse_article(article.title, article.content)
                result = await database.save_article(article, parsed)
                if result["status"] == "saved":
                    saved += 1
                elif result["status"] == "updated":
                    updated += 1
                elif result["status"] == "skipped":
                    skipped += 1
                elif result["status"] == "filtered":
                    filtered += 1

            logger.info(f"  [{name}] 抓取={len(articles)} 新增={saved} 更新={updated} 跳过={skipped} 过滤={filtered}")
        except Exception as e:
            logger.error(f"  [{name}] 错误: {e}")


async def cmd_backfill():
    """回补全文：用 Playwright 打开文章链接获取全文 + 重新 AI 解析"""
    from playwright.async_api import async_playwright
    from .scraper import AUTH_STATE_PATH
    from .content_fetcher import fetch_full_content_playwright

    articles = await database.get_articles_needing_backfill(max_content_len=500)
    logger.info(f"需要回补全文的文章: {len(articles)} 篇")

    if not articles:
        logger.info("无需回补，所有文章已有全文")
        return

    # 1. 启动 Playwright 浏览器（复用登录态）
    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=False)
        if AUTH_STATE_PATH.exists():
            context = await browser.new_context(storage_state=str(AUTH_STATE_PATH))
        else:
            context = await browser.new_context()
        page = await context.new_page()

        # 2. Playwright 逐篇获取全文
        urls = [a["url"] for a in articles]
        fetched = await fetch_full_content_playwright(page, urls, delay=1.5)

        await browser.close()

    success_count = sum(1 for v in fetched.values() if v.success)
    logger.info(f"全文获取: {success_count}/{len(urls)} 成功")

    # 3. 逐篇更新内容 + 重新 AI 解析
    updated, ai_parsed, failed = 0, 0, 0
    for art in articles:
        fc = fetched.get(art["url"])
        if not fc or not fc.success or not fc.full_text:
            failed += 1
            continue

        # 用全文重新 AI 解析
        parsed = await parse_article(art["title"], fc.full_text)
        ai_parsed += 1

        # 更新数据库
        await database.update_article_content(
            article_id=art["id"],
            content=fc.full_text,
            publish_date=fc.publish_date,
            parsed=parsed if not parsed.parse_failed else None,
        )
        updated += 1

        if updated % 20 == 0:
            logger.info(f"  AI解析进度: {updated}/{len(articles)}")

    logger.info(f"回补完成: 更新={updated} AI解析={ai_parsed} 失败={failed}")


async def cmd_stats():
    """输出统计信息"""
    stats = await database.get_stats()
    print(json.dumps(stats, ensure_ascii=False, indent=2))


def main():
    parser = argparse.ArgumentParser(description="A股主题跟踪 CLI")
    sub = parser.add_subparsers(dest="command")

    fetch_parser = sub.add_parser("fetch", help="抓取数据")
    fetch_parser.add_argument("--source", choices=list(SOURCES.keys()), help="指定数据源")

    sub.add_parser("backfill", help="回补全文（crawl4ai 拉全文 + 重新AI解析）")
    sub.add_parser("stats", help="查看统计")

    args = parser.parse_args()

    if args.command == "fetch":
        asyncio.run(cmd_fetch(args.source))
    elif args.command == "backfill":
        asyncio.run(cmd_backfill())
    elif args.command == "stats":
        asyncio.run(cmd_stats())
    else:
        parser.print_help()


if __name__ == "__main__":
    main()

"""Playwright 全文抓取模块 — 在已登录浏览器中逐篇获取文章全文 + 日期

功能：
- 复用 Playwright 浏览器会话（已登录状态）
- 逐篇打开文章 URL，提取全文 + 精确发布日期
- 自动等待 JS 渲染
- 请求间隔控制避免触发反爬
"""

import asyncio
import re
import logging
from dataclasses import dataclass
from playwright.async_api import Page, Browser

logger = logging.getLogger("a_stock_watcher.content_fetcher")


@dataclass
class FetchedContent:
    """单篇文章的全文抓取结果"""
    url: str
    full_text: str = ""
    publish_date: str = ""
    success: bool = False
    error: str = ""


def _extract_date(text: str) -> str:
    """从文本中提取发布日期"""
    match = re.search(r'(\d{4}-\d{2}-\d{2})\s*\d{2}:\d{2}:\d{2}', text)
    if match:
        return match.group(1)
    match = re.search(r'(\d{4}-\d{2}-\d{2})', text)
    if match:
        return match.group(1)
    return ""


# JS 脚本：提取文章全文 + 日期
EXTRACT_ARTICLE_JS = r"""() => {
    // 提取日期
    const allText = document.body.innerText || '';
    const dateMatch = allText.match(/(\d{4}-\d{2}-\d{2}\s+\d{2}:\d{2}:\d{2})/);
    const publishDate = dateMatch ? dateMatch[1].substring(0, 10) : '';
    
    // 提取正文：优先找文章内容容器
    let content = '';
    
    // 九阳公社文章页的内容选择器
    const selectors = [
        '.article-content',
        '.rich-text',
        '.post-content',
        '.detail-content',
        '[class*="content"]',
        '[class*="article"]',
    ];
    
    for (const sel of selectors) {
        const el = document.querySelector(sel);
        if (el && el.innerText.length > 200) {
            content = el.innerText;
            break;
        }
    }
    
    // 如果选择器没找到，取 body 文本但去掉导航
    if (!content || content.length < 200) {
        content = document.body.innerText || '';
    }
    
    return {
        content: content.substring(0, 15000),  // 限制长度
        publishDate: publishDate,
    };
}"""


async def fetch_full_content_playwright(
    page: Page,
    urls: list[str],
    delay: float = 1.5,
) -> dict[str, FetchedContent]:
    """
    在已打开的 Playwright page 中逐篇获取文章全文。

    Args:
        page: 已登录的 Playwright Page 对象
        urls: 要抓取的 URL 列表（自动去重）
        delay: 每篇之间的延迟（秒）

    Returns:
        {url: FetchedContent} 字典
    """
    # URL 去重（保持顺序）
    unique_urls = list(dict.fromkeys(urls))
    if len(unique_urls) < len(urls):
        logger.info(f"  URL 去重: {len(urls)} → {len(unique_urls)} 个唯一 URL")

    results: dict[str, FetchedContent] = {}

    for i, url in enumerate(unique_urls):
        try:
            await page.goto(url, wait_until="networkidle", timeout=15000)
            await page.wait_for_timeout(1000)  # 额外等待 JS 渲染

            data = await page.evaluate(EXTRACT_ARTICLE_JS)
            content = data.get("content", "")
            pub_date = data.get("publishDate", "")

            if content and len(content) > 100:
                results[url] = FetchedContent(
                    url=url,
                    full_text=content,
                    publish_date=pub_date,
                    success=True,
                )
            else:
                results[url] = FetchedContent(
                    url=url,
                    success=False,
                    error=f"content_len={len(content)}",
                )
        except Exception as e:
            logger.error(f"  抓取异常 {url}: {e}")
            results[url] = FetchedContent(
                url=url,
                success=False,
                error=str(e)[:100],
            )

        # 进度日志
        if (i + 1) % 5 == 0:
            ok = sum(1 for v in results.values() if v.success)
            logger.info(f"  全文抓取进度: {i+1}/{len(unique_urls)} (成功={ok})")

        # 延迟避免反爬
        if delay > 0:
            await asyncio.sleep(delay)

    return results

"""SQLite 异步数据层 — 5 表关联：articles / stocks / themes / article_stocks / article_themes"""

import aiosqlite
import json
from pathlib import Path

from .models import Article, ParsedResult, StockMention, ThemeMention

# 数据库文件位置：项目根目录/data/themes.db
DB_DIR = Path(__file__).parent.parent / "data"
DB_PATH = DB_DIR / "themes.db"

SCHEMA_SQL = """
CREATE TABLE IF NOT EXISTS articles (
    id           INTEGER PRIMARY KEY AUTOINCREMENT,
    source       TEXT NOT NULL,
    title        TEXT NOT NULL,
    url          TEXT DEFAULT '',
    content      TEXT DEFAULT '',
    images       TEXT DEFAULT '[]',
    publish_date TEXT DEFAULT '',
    fetched_at   TEXT NOT NULL,
    content_hash TEXT NOT NULL UNIQUE,
    relevance    INTEGER DEFAULT 0,
    logic_summary TEXT DEFAULT '',
    raw_yaml     TEXT DEFAULT '',
    parse_status TEXT DEFAULT 'pending'
);

CREATE TABLE IF NOT EXISTS stocks (
    id   INTEGER PRIMARY KEY AUTOINCREMENT,
    code TEXT NOT NULL UNIQUE,
    name TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS themes (
    id       INTEGER PRIMARY KEY AUTOINCREMENT,
    name     TEXT NOT NULL UNIQUE,
    category TEXT DEFAULT ''
);

CREATE TABLE IF NOT EXISTS article_stocks (
    id         INTEGER PRIMARY KEY AUTOINCREMENT,
    article_id INTEGER NOT NULL REFERENCES articles(id),
    stock_id   INTEGER NOT NULL REFERENCES stocks(id),
    context    TEXT DEFAULT '',
    logic      TEXT DEFAULT '',
    UNIQUE(article_id, stock_id)
);

CREATE TABLE IF NOT EXISTS article_themes (
    id         INTEGER PRIMARY KEY AUTOINCREMENT,
    article_id INTEGER NOT NULL REFERENCES articles(id),
    theme_id   INTEGER NOT NULL REFERENCES themes(id),
    UNIQUE(article_id, theme_id)
);

CREATE INDEX IF NOT EXISTS idx_articles_source ON articles(source);
CREATE INDEX IF NOT EXISTS idx_articles_date ON articles(publish_date);
CREATE INDEX IF NOT EXISTS idx_articles_fetched ON articles(fetched_at);
"""


async def _get_db() -> aiosqlite.Connection:
    """获取数据库连接（自动建表）"""
    DB_DIR.mkdir(parents=True, exist_ok=True)
    db = await aiosqlite.connect(str(DB_PATH))
    db.row_factory = aiosqlite.Row
    await db.executescript(SCHEMA_SQL)
    return db


# ─── 去重检查 ───────────────────────────────────────────────────────

async def check_exists(content_hash: str) -> bool:
    """检查文章是否已存在（通过 content_hash）。用于 AI 解析之前去重。"""
    db = await _get_db()
    try:
        cursor = await db.execute(
            "SELECT 1 FROM articles WHERE content_hash = ?", (content_hash,)
        )
        return (await cursor.fetchone()) is not None
    finally:
        await db.close()


# ─── 写入 ─────────────────────────────────────────────────────────

async def save_article(article: Article, parsed: ParsedResult | None = None) -> dict:
    """
    保存文章 + AI 解析结果（自动去重、自动创建 stock/theme 实体）。
    返回 {"status": "saved"|"skipped"|"filtered", "article_id": int|None}
    """
    # 噪音过滤
    if parsed and parsed.relevance < 5:
        return {"status": "filtered", "article_id": None,
                "reason": parsed.filter_reason or f"relevance={parsed.relevance}"}

    db = await _get_db()
    try:
        # 插入文章（去重）
        try:
            cursor = await db.execute(
                """INSERT INTO articles
                   (source, title, url, content, images, publish_date, fetched_at,
                    content_hash, relevance, logic_summary, raw_yaml, parse_status)
                   VALUES (?,?,?,?,?,?,?,?,?,?,?,?)""",
                (
                    article.source, article.title, article.url, article.content,
                    json.dumps(article.images, ensure_ascii=False),
                    parsed.publish_date if parsed else article.publish_date,
                    article.fetched_at, article.content_hash,
                    parsed.relevance if parsed else 0,
                    parsed.logic_summary if parsed else "",
                    parsed.raw_yaml if parsed else "",
                    "parsed" if parsed and not parsed.parse_failed else "pending",
                ),
            )
            article_id = cursor.lastrowid
        except aiosqlite.IntegrityError:
            # 产业链来源：同一 industry_id 有新版本 → 覆盖旧数据（保留最新）
            if article.source == "industry_chain":
                cur2 = await db.execute(
                    "SELECT id FROM articles WHERE content_hash=?", (article.content_hash,)
                )
                row = await cur2.fetchone()
                if not row:
                    return {"status": "skipped", "article_id": None}
                article_id = row["id"]
                await db.execute(
                    """UPDATE articles
                       SET title=?, content=?, images=?, publish_date=?, fetched_at=?,
                           relevance=?, logic_summary=?, raw_yaml=?, parse_status=?
                       WHERE id=?""",
                    (
                        article.title, article.content,
                        json.dumps(article.images, ensure_ascii=False),
                        parsed.publish_date if parsed else article.publish_date,
                        article.fetched_at,
                        parsed.relevance if parsed else 0,
                        parsed.logic_summary if parsed else "",
                        parsed.raw_yaml if parsed else "",
                        "parsed" if parsed and not parsed.parse_failed else "pending",
                        article_id,
                    ),
                )
                # 重建关联数据
                if parsed and not parsed.parse_failed:
                    await db.execute("DELETE FROM article_stocks WHERE article_id=?", (article_id,))
                    await db.execute("DELETE FROM article_themes WHERE article_id=?", (article_id,))
                    for s in parsed.stocks:
                        stock_id = await _upsert_stock(db, s.code, s.name)
                        await _safe_insert(
                            db,
                            "INSERT OR IGNORE INTO article_stocks (article_id, stock_id, context, logic) VALUES (?,?,?,?)",
                            (article_id, stock_id, s.context, s.logic),
                        )
                    for t in parsed.themes:
                        theme_id = await _upsert_theme(db, t.name, t.category)
                        await _safe_insert(
                            db,
                            "INSERT OR IGNORE INTO article_themes (article_id, theme_id) VALUES (?,?)",
                            (article_id, theme_id),
                        )
                await db.commit()
                return {"status": "updated", "article_id": article_id}
            return {"status": "skipped", "article_id": None}

        # 写入关联数据
        if parsed and not parsed.parse_failed:
            for s in parsed.stocks:
                stock_id = await _upsert_stock(db, s.code, s.name)
                await _safe_insert(
                    db,
                    "INSERT OR IGNORE INTO article_stocks (article_id, stock_id, context, logic) VALUES (?,?,?,?)",
                    (article_id, stock_id, s.context, s.logic),
                )
            for t in parsed.themes:
                theme_id = await _upsert_theme(db, t.name, t.category)
                await _safe_insert(
                    db,
                    "INSERT OR IGNORE INTO article_themes (article_id, theme_id) VALUES (?,?)",
                    (article_id, theme_id),
                )

        await db.commit()
        return {"status": "saved", "article_id": article_id}
    finally:
        await db.close()


async def _upsert_stock(db: aiosqlite.Connection, code: str, name: str) -> int:
    """插入或获取股票 ID"""
    try:
        cursor = await db.execute(
            "INSERT INTO stocks (code, name) VALUES (?, ?)", (code, name)
        )
        return cursor.lastrowid
    except aiosqlite.IntegrityError:
        cursor = await db.execute("SELECT id FROM stocks WHERE code = ?", (code,))
        row = await cursor.fetchone()
        return row["id"]


async def _upsert_theme(db: aiosqlite.Connection, name: str, category: str) -> int:
    """插入或获取主题 ID"""
    try:
        cursor = await db.execute(
            "INSERT INTO themes (name, category) VALUES (?, ?)", (name, category)
        )
        return cursor.lastrowid
    except aiosqlite.IntegrityError:
        cursor = await db.execute("SELECT id FROM themes WHERE name = ?", (name,))
        row = await cursor.fetchone()
        return row["id"]


async def _safe_insert(db: aiosqlite.Connection, sql: str, params: tuple):
    """安全插入，忽略重复"""
    try:
        await db.execute(sql, params)
    except aiosqlite.IntegrityError:
        pass


# ─── 查询 ─────────────────────────────────────────────────────────

async def query_by_stock(stock: str, limit: int = 50) -> list[dict]:
    """按股票代码或名称查询相关文章（时间倒序）"""
    db = await _get_db()
    try:
        cursor = await db.execute(
            """SELECT a.*, s.code as stock_code, s.name as stock_name,
                      ars.context, ars.logic
               FROM articles a
               JOIN article_stocks ars ON a.id = ars.article_id
               JOIN stocks s ON s.id = ars.stock_id
               WHERE s.code = ? OR s.name LIKE ?
               ORDER BY a.publish_date DESC, a.fetched_at DESC
               LIMIT ?""",
            (stock, f"%{stock}%", limit),
        )
        return [_row_to_dict(r) for r in await cursor.fetchall()]
    finally:
        await db.close()


async def query_by_theme(theme: str, limit: int = 50) -> list[dict]:
    """按主题名称查询关联文章和股票"""
    db = await _get_db()
    try:
        cursor = await db.execute(
            """SELECT a.*, t.name as theme_name, t.category as theme_category
               FROM articles a
               JOIN article_themes at2 ON a.id = at2.article_id
               JOIN themes t ON t.id = at2.theme_id
               WHERE t.name LIKE ?
               ORDER BY a.publish_date DESC, a.fetched_at DESC
               LIMIT ?""",
            (f"%{theme}%", limit),
        )
        articles = [_row_to_dict(r) for r in await cursor.fetchall()]

        # 附加每篇文章的关联股票
        for art in articles:
            cursor2 = await db.execute(
                """SELECT s.code, s.name, ars.context, ars.logic
                   FROM article_stocks ars JOIN stocks s ON s.id = ars.stock_id
                   WHERE ars.article_id = ?""",
                (art["id"],),
            )
            art["stocks"] = [dict(r) for r in await cursor2.fetchall()]
        return articles
    finally:
        await db.close()


async def query_latest(days: int = 7, source: str | None = None, limit: int = 50) -> list[dict]:
    """查询最近 N 天的新增文章"""
    db = await _get_db()
    try:
        conditions = ["fetched_at >= datetime('now', ?)"]
        params: list = [f"-{days} days"]
        if source:
            conditions.append("source = ?")
            params.append(source)
        where = "WHERE " + " AND ".join(conditions)
        params.append(limit)

        cursor = await db.execute(
            f"SELECT * FROM articles {where} ORDER BY fetched_at DESC LIMIT ?",
            params,
        )
        return [_row_to_dict(r) for r in await cursor.fetchall()]
    finally:
        await db.close()


async def get_stock_timeline(stock: str) -> list[dict]:
    """获取某只股票的逻辑演变时间线"""
    db = await _get_db()
    try:
        cursor = await db.execute(
            """SELECT a.publish_date, a.title, a.source, a.url,
                      ars.context, ars.logic, a.logic_summary
               FROM articles a
               JOIN article_stocks ars ON a.id = ars.article_id
               JOIN stocks s ON s.id = ars.stock_id
               WHERE s.code = ? OR s.name LIKE ?
               ORDER BY a.publish_date ASC""",
            (stock, f"%{stock}%"),
        )
        return [dict(r) for r in await cursor.fetchall()]
    finally:
        await db.close()


async def get_stats() -> dict:
    """获取各数据源的统计信息"""
    db = await _get_db()
    try:
        cursor = await db.execute(
            """SELECT source, COUNT(*) as count, MAX(fetched_at) as last_fetch
               FROM articles GROUP BY source"""
        )
        rows = await cursor.fetchall()
        stats = {
            "sources": {row["source"]: {"count": row["count"], "last_fetch": row["last_fetch"]} for row in rows},
        }
        # 总计
        cursor = await db.execute("SELECT COUNT(*) as c FROM stocks")
        stats["total_stocks"] = (await cursor.fetchone())["c"]
        cursor = await db.execute("SELECT COUNT(*) as c FROM themes")
        stats["total_themes"] = (await cursor.fetchone())["c"]
        return stats
    finally:
        await db.close()


def _row_to_dict(row) -> dict:
    """将 Row 对象转为 dict，解析 JSON 字段"""
    d = dict(row)
    for json_field in ("images",):
        if json_field in d and isinstance(d[json_field], str):
            try:
                d[json_field] = json.loads(d[json_field])
            except (json.JSONDecodeError, TypeError):
                pass
    return d


# ─── Backfill（回补全文+重新解析）───────────────────────────────────

async def get_articles_needing_backfill(max_content_len: int = 500) -> list[dict]:
    """
    获取需要回补全文的文章（内容过短 + 有 URL 的）。
    返回 [{id, url, title, source}, ...]，URL 去重。
    """
    db = await _get_db()
    try:
        cursor = await db.execute(
            """SELECT id, url, title, source FROM articles
               WHERE length(content) <= ? AND url != '' AND url IS NOT NULL
               ORDER BY id""",
            (max_content_len,),
        )
        rows = await cursor.fetchall()
        # URL 去重：同一 URL 只保留第一个
        seen_urls = set()
        results = []
        for row in rows:
            url = row["url"]
            if url not in seen_urls:
                seen_urls.add(url)
                results.append(dict(row))
        return results
    finally:
        await db.close()


async def update_article_content(
    article_id: int,
    content: str,
    publish_date: str = "",
    parsed: "ParsedResult | None" = None,
) -> None:
    """
    更新文章的内容和日期，并可选重新写入 AI 解析结果。
    如果提供了 parsed，会清除旧的 article_stocks/article_themes 并重建。
    """
    db = await _get_db()
    try:
        # 更新文章基本字段
        updates = {"content": content}
        if publish_date:
            updates["publish_date"] = publish_date
        if parsed and not parsed.parse_failed:
            updates["relevance"] = parsed.relevance
            updates["logic_summary"] = parsed.logic_summary
            updates["raw_yaml"] = parsed.raw_yaml
            updates["parse_status"] = "parsed"
            if parsed.publish_date:
                updates["publish_date"] = parsed.publish_date

        set_clause = ", ".join(f"{k}=?" for k in updates)
        values = list(updates.values()) + [article_id]
        await db.execute(f"UPDATE articles SET {set_clause} WHERE id=?", values)

        # 重建关联数据
        if parsed and not parsed.parse_failed:
            await db.execute("DELETE FROM article_stocks WHERE article_id=?", (article_id,))
            await db.execute("DELETE FROM article_themes WHERE article_id=?", (article_id,))

            for s in parsed.stocks:
                stock_id = await _upsert_stock(db, s.code, s.name)
                await _safe_insert(
                    db,
                    "INSERT OR IGNORE INTO article_stocks (article_id, stock_id, context, logic) VALUES (?,?,?,?)",
                    (article_id, stock_id, s.context, s.logic),
                )
            for t in parsed.themes:
                theme_id = await _upsert_theme(db, t.name, t.category)
                await _safe_insert(
                    db,
                    "INSERT OR IGNORE INTO article_themes (article_id, theme_id) VALUES (?,?)",
                    (article_id, theme_id),
                )

        await db.commit()
    finally:
        await db.close()

"""数据模型定义"""

from dataclasses import dataclass, field
from datetime import datetime
import hashlib


@dataclass
class Article:
    """原始文章条目"""
    title: str
    source: str  # study_hot | industry_chain | action
    url: str = ""
    content: str = ""  # 全文内容
    images: list[str] = field(default_factory=list)  # 图片路径列表
    publish_date: str = ""  # 发布日期
    fetched_at: str = field(default_factory=lambda: datetime.now().isoformat())
    content_hash: str = ""
    _parsed: object = field(default=None, repr=False)  # 已有的AI解析结果

    def __post_init__(self):
        if not self.content_hash:
            raw = f"{self.source}:{self.title}:{self.url}"
            self.content_hash = hashlib.sha256(raw.encode()).hexdigest()[:16]


@dataclass
class StockMention:
    """文章中提及的股票"""
    code: str  # 股票代码，如 600519
    name: str  # 股票名称，如 贵州茅台
    context: str = ""  # 提及的上下文片段
    logic: str = ""  # 投资逻辑


@dataclass
class ThemeMention:
    """文章中涉及的主题"""
    name: str  # 主题名称，如 "白酒复苏"
    category: str = ""  # 分类，如 "消费"、"科技"


@dataclass
class ParsedResult:
    """AI 解析文章后的结构化结果"""
    relevance: int = 0  # 0-10，<5 自动丢弃
    publish_date: str = ""
    stocks: list[StockMention] = field(default_factory=list)
    themes: list[ThemeMention] = field(default_factory=list)
    logic_summary: str = ""
    raw_yaml: str = ""  # AI 原始输出
    parse_failed: bool = False
    filter_reason: str = ""  # 被过滤的原因

"""Playwright 爬虫基类 — 管理浏览器生命周期 + 登录态持久化 + 登录弹窗容错"""

import logging
from abc import ABC, abstractmethod
from pathlib import Path
from playwright.async_api import async_playwright, Page, BrowserContext

from .models import Article

logger = logging.getLogger("a_stock_watcher.scraper")

# 登录态存储路径
AUTH_DIR = Path(__file__).parent.parent / "data"
AUTH_STATE_PATH = AUTH_DIR / "auth_state.json"


class BaseScraper(ABC):
    """
    爬虫基类。子类只需实现 parse() 方法。
    
    登录态持久化：
        - 首次登录后，cookie/localStorage 保存到 data/auth_state.json
        - 后续运行自动加载，无需重复登录
        - 如果登录过期，运行 `uv run python -m a_stock_watcher.auth` 重新登录
    """

    # 子类覆盖
    source_name: str = ""
    source_url: str = ""

    def __init__(self, headless: bool = True):
        """
        Args:
            headless: 是否无头模式。默认 True（无头），调试时设 False。
        """
        self.headless = headless

    async def run(self) -> list[Article]:
        """启动浏览器 → 加载登录态 → 导航到目标页面 → 解析 → 关闭浏览器"""
        async with async_playwright() as p:
            browser = await p.chromium.launch(headless=self.headless)
            try:
                # 加载已保存的登录态
                context = await self._create_context(browser)
                page = await context.new_page()
                await page.goto(self.source_url, wait_until="networkidle", timeout=60000)
                await self.dismiss_login_modal(page)
                articles = await self.parse(page)
                return articles
            finally:
                await browser.close()

    async def _create_context(self, browser) -> BrowserContext:
        """创建带登录态的浏览器上下文"""
        if AUTH_STATE_PATH.exists():
            return await browser.new_context(storage_state=str(AUTH_STATE_PATH))
        return await browser.new_context()

    async def dismiss_login_modal(self, page: Page):
        """关闭登录弹窗（通用容错）

        韭菜公社页面偶尔弹出登录框，需要点击右上角叉号关闭。
        策略：先尝试点击关闭按钮，再用 JS 暴力删除弹窗 DOM。
        """
        try:
            # 1. 尝试点击关闭按钮（右上角叉号）
            close_selectors = [
                '.el-dialog__headerbtn',        # Element UI 弹窗关闭按钮
                '.el-icon-close',               # Element UI 关闭图标
                '.close-btn',                   # 通用关闭按钮
                'button[aria-label="Close"]',    # 无障碍标签
            ]
            for selector in close_selectors:
                btn = await page.query_selector(selector)
                if btn and await btn.is_visible():
                    await btn.click()
                    await page.wait_for_timeout(500)
                    logger.info("  关闭登录弹窗（点击关闭按钮）")
                    return

            # 2. 回退：JS 删除弹窗 DOM
            removed = await page.evaluate("""() => {
                let removed = 0;
                document.querySelectorAll('.el-dialog__wrapper, .v-modal, .login-modal').forEach(el => {
                    el.remove();
                    removed++;
                });
                document.body.classList.remove('el-popup-parent--hidden');
                return removed;
            }""")
            if removed:
                logger.info(f"  关闭登录弹窗（JS删除 {removed} 个元素）")
        except Exception:
            pass  # 弹窗处理失败不影响主流程

    @abstractmethod
    async def parse(self, page: Page) -> list[Article]:
        """
        子类实现：从页面中提取文章数据。
        
        Args:
            page: 已导航到 source_url 的 Playwright Page 对象
            
        Returns:
            解析出的 Article 列表（含 title, content, images, url 等）
        """
        ...

"""
A股主题跟踪 MCP Server

通过 FastMCP 创建 MCP 服务器，提供主题抓取和多维查询工具。
定时抓取由外部 cron 调度（见 cli.py）。
启动方式：python -m a_stock_watcher
"""

from fastmcp import FastMCP

from .sources import SOURCES
from . import database
from .ai_parser import parse_article

# 创建 MCP 服务器实例
mcp = FastMCP(
    "A-Stock Watcher",
    instructions=(
        "A股主题跟踪工具。从九阳公社等平台抓取热点研究、产业链、异动等主题数据。"
        "使用 Gemini 2.5 Flash AI 自动提取股票、主题、投资逻辑，过滤噪音。"
        "支持按股票、主题、时间等多维度查询和溯源。"
    ),
)


# ─── 数据抓取 ───────────────────────────────────────────────────

@mcp.tool()
async def fetch_themes(source: str | None = None) -> dict:
    """
    抓取指定来源的主题数据 → AI 解析 → 入库（自动去重+噪音过滤）。

    Args:
        source: 数据源名称（study_hot / industry_chain / action）。
               不指定则抓取全部来源。
    """
    targets = {}
    if source:
        if source not in SOURCES:
            return {"error": f"未知数据源: {source}，可用: {list(SOURCES.keys())}"}
        targets[source] = SOURCES[source]
    else:
        targets = SOURCES

    results = {}
    for name, scraper_cls in targets.items():
        try:
            scraper = scraper_cls()
            articles = await scraper.run()
            saved, skipped, filtered = 0, 0, 0

            for article in articles:
                # 去重前置：先查DB，已存在的跳过不调AI
                if await database.check_exists(article.content_hash):
                    skipped += 1
                    continue

                # 如果已有解析结果（如 industry_chain 的图片解析），直接使用
                if article._parsed is not None:
                    parsed = article._parsed
                else:
                    parsed = await parse_article(article.title, article.content)
                result = await database.save_article(article, parsed)
                if result["status"] == "saved":
                    saved += 1
                elif result["status"] == "skipped":
                    skipped += 1
                elif result["status"] == "filtered":
                    filtered += 1

            results[name] = {
                "fetched": len(articles),
                "saved": saved,
                "skipped": skipped,
                "filtered": filtered,
            }
        except Exception as e:
            results[name] = {"error": str(e)}

    return results


@mcp.tool()
async def fetch_all() -> dict:
    """一次性抓取所有来源 → AI 解析 → 入库。"""
    return await fetch_themes()


# ─── 多维查询 ───────────────────────────────────────────────────

@mcp.tool()
async def list_sources() -> dict:
    """列出所有数据源及统计信息（数据量、最近抓取时间、股票/主题总数）。"""
    stats = await database.get_stats()
    sources_info = {}
    for name, scraper_cls in SOURCES.items():
        source_stats = stats.get("sources", {}).get(name, {})
        sources_info[name] = {
            "url": scraper_cls.source_url,
            "total_items": source_stats.get("count", 0),
            "last_fetch": source_stats.get("last_fetch", "从未抓取"),
        }
    sources_info["_summary"] = {
        "total_stocks": stats.get("total_stocks", 0),
        "total_themes": stats.get("total_themes", 0),
    }
    return sources_info


@mcp.tool()
async def query_by_stock(stock: str, limit: int = 50) -> list[dict]:
    """
    按股票查询相关文章和投资逻辑（时间倒序）。

    Args:
        stock: 股票代码（如 600519）或名称（如 贵州茅台）
        limit: 返回数量上限
    """
    return await database.query_by_stock(stock=stock, limit=limit)


@mcp.tool()
async def query_by_theme(theme: str, limit: int = 50) -> list[dict]:
    """
    按主题查询关联文章和股票。

    Args:
        theme: 主题名称（如 "白酒复苏"、"AI"）
        limit: 返回数量上限
    """
    return await database.query_by_theme(theme=theme, limit=limit)


@mcp.tool()
async def query_latest(days: int = 7, source: str | None = None, limit: int = 50) -> list[dict]:
    """
    查询最近 N 天的新增文章。

    Args:
        days: 天数范围，默认 7 天
        source: 可选，按数据源筛选
        limit: 返回数量上限
    """
    return await database.query_latest(days=days, source=source, limit=limit)


@mcp.tool()
async def get_stock_timeline(stock: str) -> list[dict]:
    """
    获取某只股票的逻辑演变时间线（按日期升序）。
    可用于溯源：看某公司的逻辑、段子是何时首次出现的。

    Args:
        stock: 股票代码或名称
    """
    return await database.get_stock_timeline(stock=stock)



def main() -> None:
    """启动 MCP 服务器（stdio 模式）"""
    mcp.run()


if __name__ == "__main__":
    main()

"""数据源注册表"""

from .study_hot import StudyHotScraper
from .industry_chain import IndustryChainScraper
from .action import ActionScraper

# source_name → Scraper 类的映射
import functools

# source_name → Scraper 工厂（带默认参数）
SOURCES = {
    "study_hot": StudyHotScraper,                                         # scroll_rounds=5 → limit=25
    "industry_chain": IndustryChainScraper,                               # max_items=15
    "action": functools.partial(ActionScraper, lookback_days=2),          # 日常增量只拉近 2 天
}

__all__ = ["SOURCES", "StudyHotScraper", "IndustryChainScraper", "ActionScraper"]

"""异动 — https://www.jiuyangongshe.com/action

实现方式（B计划 API拦截）：
- 拦截页面发出的 POST /v1/action/field 请求（带 date 参数）
- 一次请求返回当天全部板块 + 每板块内所有个股的异动解析
- 无需逐天导航 + 点击展开 DOM

API 响应结构：
    data: [{
        name,       # 板块名，如"电力"、"算力"
        reason,     # 板块整体解析
        date,       # 日期
        list: [{
            code,   # 股票代码，如 sh688525
            name,   # 股票名称
            article: {
                title,
                action_info: {
                    expound,  # 异动解析全文
                    price, shares_range, time, ...
                }
            }
        }]
    }]
"""

import json
import logging
from datetime import datetime, timedelta
from playwright.async_api import Page

from ..scraper import BaseScraper
from ..models import Article

logger = logging.getLogger("a_stock_watcher.action")


class ActionScraper(BaseScraper):
    source_name = "action"
    source_url = "https://www.jiuyangongshe.com/action"

    def __init__(self, headless: bool = False, lookback_days: int = 6):
        super().__init__(headless=headless)
        self.lookback_days = lookback_days

    async def parse(self, page: Page) -> list[Article]:
        today = datetime.now()
        dates = [
            (today - timedelta(days=i)).strftime("%Y-%m-%d")
            for i in range(self.lookback_days)
        ]

        articles: list[Article] = []
        for date_str in dates:
            day_articles = await self._fetch_day(page, date_str)
            articles.extend(day_articles)
            logger.info(f"  action [{date_str}] 提取 {len(day_articles)} 个板块")

        return articles

    async def _fetch_day(self, page: Page, date_str: str) -> list[Article]:
        """拦截 action/field API，获取单日全部板块异动数据"""
        captured: list[dict] = []

        async def on_response(resp):
            if "/action/field" in resp.url and resp.request.resource_type in ("xhr", "fetch"):
                try:
                    body = await resp.json()
                    captured.append(body)
                except Exception:
                    pass

        page.on("response", on_response)
        await page.goto(
            f"{self.source_url}/{date_str}",
            wait_until="networkidle",
            timeout=30000,
        )
        await page.wait_for_timeout(2000)

        # 点"全部异动解析" tab 触发 action/field 请求
        tab = await page.query_selector('text="全部异动解析"')
        if tab:
            await tab.click()
            await page.wait_for_timeout(2000)

        page.remove_listener("response", on_response)

        if not captured:
            logger.warning(f"  [{date_str}] 未捕获 action/field 响应")
            return []

        # 取数据最多的响应
        best = max(captured, key=lambda b: len(b.get("data") or []))
        fields: list[dict] = best.get("data") or []

        articles: list[Article] = []
        for field in fields:
            field_name = (field.get("name") or "").strip()
            # 跳过"简图"（纯图表，无股票列表）
            if not field_name or field_name == "简图":
                continue

            stocks = field.get("list") or []
            if not stocks:
                continue

            # 每个板块聚合成一篇 Article
            lines: list[str] = []
            reason = (field.get("reason") or "").strip()
            if reason:
                lines.append(f"板块逻辑：{reason}\n")

            for stock in stocks:
                code = (stock.get("code") or "").strip()
                name = (stock.get("name") or "").strip()
                article = stock.get("article") or {}
                action_info = article.get("action_info") or {}
                expound = (action_info.get("expound") or "").strip()
                price = action_info.get("price")
                shares = action_info.get("shares_range")

                meta = ""
                if price:
                    meta += f" 价格={price}"
                if shares:
                    meta += f" 异动量={shares}万"

                if name:
                    line = f"【{name} {code}】{meta}"
                    if expound:
                        line += f"\n{expound}"
                    lines.append(line)

            content = "\n\n".join(lines)
            if not content.strip():
                continue

            articles.append(Article(
                title=f"[{date_str}异动] {field_name}（{len(stocks)}股）",
                source=self.source_name,
                url=f"{self.source_url}/{date_str}",
                content=content,
                images=[],
                publish_date=date_str,
            ))

        return articles

"""产业链 — https://www.jiuyangongshe.com/industryChain

实现方式（B计划 API拦截）：
- 拦截页面发出的 POST /v1/industry/list 请求
- 直接从 API JSON 响应提取产业链结构化数据
- 无需点击详情页、截图、Gemini Vision 识别
- 速度从 ~5min → ~30s，且去掉了 Gemini 图像识别依赖

去重策略：
- content_hash 基于 industry_id（UUID），与 title 里的日期无关
- 同一产业链有新版本时，数据库做 upsert 保留最新内容

API 响应结构：
    data.result: [{
        industry_id, title, keyword, content,
        imgs: [...],  # 产业链信息图 URL 列表
    }]
"""

import hashlib
import json
import logging
from playwright.async_api import Page

from ..scraper import BaseScraper
from ..models import Article

logger = logging.getLogger("a_stock_watcher.industry_chain")


class IndustryChainScraper(BaseScraper):
    source_name = "industry_chain"
    source_url = "https://www.jiuyangongshe.com/industryChain"

    def __init__(self, headless: bool = False, max_items: int = 15):
        """
        Args:
            headless: 无头模式
            max_items: 最多返回的产业链条数（API 分页，pageSize 传入）
        """
        super().__init__(headless=headless)
        self.max_items = max_items

    async def parse(self, page: Page) -> list[Article]:
        captured: list[dict] = []

        # 只监听响应，不修改请求（修改 body 会导致 token 校验失败）
        async def on_response(resp):
            if "/industry/list" in resp.url and resp.request.resource_type in ("xhr", "fetch"):
                try:
                    body = await resp.json()
                    captured.append(body)
                except Exception:
                    pass

        page.on("response", on_response)

        # 重新加载页面，触发 industry/list API 请求
        await page.reload(wait_until="networkidle", timeout=30000)
        await page.wait_for_timeout(2000)

        page.remove_listener("response", on_response)

        if not captured:
            logger.error("未捕获 industry/list API 响应，返回空列表")
            return []

        # 取数据最多的那次响应
        best = max(captured, key=lambda b: len((b.get("data") or {}).get("result", [])))
        result: list[dict] = (best.get("data") or {}).get("result", [])
        logger.info(f"industry/list API 返回 {len(result)} 条产业链")

        articles: list[Article] = []
        for rank, item in enumerate(result[:self.max_items], start=1):
            title = (item.get("title") or "").strip()
            industry_id = str(item.get("industry_id") or "")
            if not title:
                continue

            # content 字段：产业链文字描述（替代 Gemini Vision OCR）
            content = (item.get("content") or "").strip()
            keyword = (item.get("keyword") or "").strip()
            if keyword and keyword not in content:
                content = f"关键词：{keyword}\n\n{content}" if content else f"关键词：{keyword}"

            # imgs 字段：产业链信息图 URL（JSON 字符串，需要 parse）
            raw_imgs = item.get("imgs") or "[]"
            try:
                imgs: list[str] = json.loads(raw_imgs) if isinstance(raw_imgs, str) else raw_imgs
            except Exception:
                imgs = []

            article = Article(
                title=f"[产业链#{rank}] {title}",
                source=self.source_name,
                url=f"{self.source_url}/{industry_id}",  # URL 含 ID，便于查询
                content=content or title,
                images=imgs,
                publish_date="",
            )
            # 用 industry_id 建稳定 hash：同一产业链不同版本 hash 相同 → upsert 保留最新
            article.content_hash = hashlib.sha256(
                f"industry_chain:{industry_id}".encode()
            ).hexdigest()[:16]
            articles.append(article)

        logger.info(f"共提取 {len(articles)} 条产业链")
        return articles

"""热点研究 — https://www.jiuyangongshe.com/study_hot

实现方式（B计划 API拦截）：
- 拦截页面发出的 POST /v1/timeline/news 请求，注入 limit 参数
- 直接从 API JSON 响应提取文章列表和全文内容
- 无需逐篇打开链接，速度从 5-15min → ~30s

API 响应结构：
    data: [{date, list: [{article_id, title, content, timeline, user: {username}}]}]
"""

import json
import logging
from playwright.async_api import Page

from ..scraper import BaseScraper
from ..models import Article

logger = logging.getLogger("a_stock_watcher.sources.study_hot")


class StudyHotScraper(BaseScraper):
    source_name = "study_hot"
    source_url = "https://www.jiuyangongshe.com/study_hot"

    def __init__(self, headless: bool = False, scroll_rounds: int = 5):
        """
        Args:
            headless: 无头模式
            scroll_rounds: 兼容旧接口，映射为 limit（每轮约 5 条，最多 30）
        """
        super().__init__(headless=headless)
        # scroll_rounds 保持接口兼容，映射到 API limit（最大 30）
        self.limit = min(scroll_rounds * 5, 30)

    async def parse(self, page: Page) -> list[Article]:
        captured: list[dict] = []

        # 拦截 timeline/news 请求，注入 limit 参数
        async def intercept(route):
            if "/timeline/news" in route.request.url:
                await route.continue_(
                    post_data=json.dumps({"limit": self.limit}),
                )
            else:
                await route.continue_()

        async def on_response(resp):
            if "timeline/news" in resp.url and resp.request.resource_type in ("xhr", "fetch"):
                try:
                    body = await resp.json()
                    captured.append(body)
                except Exception:
                    pass

        await page.route("**/*", intercept)
        page.on("response", on_response)

        # 重新加载页面，触发 timeline/news API 请求
        await page.reload(wait_until="networkidle", timeout=30000)
        await page.wait_for_timeout(2000)  # 等待响应处理完成

        await page.unroute("**/*", intercept)
        page.remove_listener("response", on_response)

        if not captured:
            logger.error("未捕获 timeline/news API 响应，返回空列表")
            return []

        # 取最后一次响应（可能多次触发，取数据最多的那次）
        news_body = max(captured, key=lambda b: len(b.get("data", [])))
        days: list[dict] = news_body.get("data", [])
        logger.info(f"timeline/news API 返回 {len(days)} 天数据")

        articles: list[Article] = []
        for day in days:
            for item in day.get("list", []):
                article_id = str(item.get("article_id", ""))
                title = (item.get("title") or "").strip()
                if not title or len(title) < 3:
                    continue

                content = (item.get("content") or "").strip() or title
                # date 字段是当天日期（timeline 是嵌套结构，不是日期字符串）
                pub_date = str(day.get("date") or "")[:10]
                url = (
                    f"https://www.jiuyangongshe.com/a/{article_id}"
                    if article_id else self.source_url
                )

                articles.append(Article(
                    title=title,
                    source=self.source_name,
                    url=url,
                    content=content,
                    images=[],
                    publish_date=pub_date,
                ))

        logger.info(f"共提取 {len(articles)} 篇文章")
        return articles

[project]
name = "a-stock-watcher"
version = "0.2.0"
description = "A股主题跟踪 MCP Server - 从九阳公社等平台抓取热点、产业链、异动数据"
requires-python = ">=3.11"
dependencies = [
    "fastmcp>=2.0.0",
    "playwright>=1.40.0",
    "aiosqlite>=0.20.0",
    "google-genai>=1.0.0",
    "pyyaml>=6.0",
]

[project.optional-dependencies]
dev = [
    "pytest>=8.0.0",
    "pytest-asyncio>=0.24.0",
]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["a_stock_watcher"]