"""Use DeepSeek to summarize AI/quant research RSS into categorized briefs.
v2: back-fills pub_time from raw RSS data after AI summarization.
"""
import json, os, requests, re
from datetime import datetime, timezone, timedelta
from dateutil import parser as dtparser

API_KEY = "sk-bbca4a0380d549389f0d27cdea0b5228"
API_URL = "https://api.deepseek.com/v1/chat/completions"
MODEL = "deepseek-v4-flash"

INPUT = "/var/www/nav/data/research.json"
OUTPUT = "/var/www/nav/data/research_summary.json"
tz = timezone(timedelta(hours=8))


def load_raw():
    if not os.path.exists(INPUT):
        return []
    with open(INPUT) as f:
        data = json.load(f)
    return data.get("items", [])


def build_time_lookup(raw_items):
    """Build lookup: title -> pub_time (Beijing time string)."""
    lookup = {}
    for item in raw_items:
        title = item.get("title", "").strip()
        published = item.get("published", "").strip()
        if not title or not published:
            continue
        try:
            dt = dtparser.parse(published)
            if dt.tzinfo:
                dt = dt.astimezone(tz)
            else:
                dt = dt.replace(tzinfo=tz)
            now = datetime.now(tz)
            if dt.date() == now.date():
                time_str = dt.strftime("%H:%M")
            else:
                time_str = dt.strftime("%m-%d %H:%M")
            lookup[title] = time_str
        except Exception:
            continue
    return lookup


def match_pub_time(items, lookup):
    for item in items:
        if "pub_time" in item:
            continue
        title = item.get("title", "").strip()
        if title in lookup:
            item["pub_time"] = lookup[title]
        else:
            for raw_title, time_val in lookup.items():
                short = min(len(raw_title), len(title))
                if short > 10 and (raw_title[:short] in title or title[:short] in raw_title):
                    item["pub_time"] = time_val
                    break


def summarize(items):
    news_text = ""
    for i, item in enumerate(items[:100], 1):
        title = item.get("title", "")
        summary = item.get("summary", "")[:300].replace("<", "").replace(">", "")
        source = item.get("source", "")
        link = item.get("link", "")
        news_text += f"{i}. [{source}] {title}\n   {summary}\n   url: {link}\n\n"

    prompt = f"""你是一个 AI/量化研究资讯编辑。请从原始 RSS 中精选出**最有价值的论文和博客文章**，按主题分类。

## 选稿标准
- ArXiv 论文：AI 重大突破（LLM、Agent、多模态）、量化交易新方法
- 博客文章：有实操价值的策略思路、代码实现、回测分析
- 过滤：纯新闻通告、产品广告、重复性内容

## 分类（每个分类 8-15 条）
- AI 前沿论文（ArXiv cs.AI / cs.LG 重要论文）
- 量化交易研究（策略思路、回测方法、风险管理）
- 行业动态（AI 公司、产品发布、监管动态）

## 输出格式
每条：rank, subtitle（15-30字）, title（必须与原始标题完全一致，这是时间匹配的关键）, source, link, summary（200-400字，说清楚这篇讲了什么、有什么用）

## 原始内容
{news_text}

JSON格式: {{"categories": [{{"name":"分类名","items":[{{"rank":1,"subtitle":"...","title":"...","source":"...","link":"...","summary":"..."}}]}}]}}"""

    resp = requests.post(
        API_URL,
        headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"},
        json={
            "model": MODEL,
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.3,
            "max_tokens": 16000,
        },
        timeout=180,
    )

    if resp.status_code != 200:
        print(f"API error: {resp.status_code}")
        return None

    content = resp.json()["choices"][0]["message"]["content"].strip()
    if content.startswith("```"):
        lines = content.split("\n")
        content = "\n".join(lines[1:-1]) if lines[-1].strip() == "```" else "\n".join(lines[1:])

    try:
        result = json.loads(content)
    except:
        content = re.sub(r',\s*}', '}', content)
        content = re.sub(r',\s*]', ']', content)
        try:
            result = json.loads(content)
        except:
            print(f"JSON parse failed, len={len(content)}")
            return None

    return result


def run():
    now = datetime.now(tz).strftime("%Y-%m-%d %H:%M")
    raw_items = load_raw()
    if not raw_items:
        print("No research data")
        exit(0)

    print(f"Loaded {len(raw_items)} research items. AI summarizing...")
    data = summarize(raw_items)

    if not data:
        print("Summarization failed")
        exit(1)

    # Back-fill pub_time
    time_lookup = build_time_lookup(raw_items)
    matched = 0
    for cat in data.get("categories", []):
        match_pub_time(cat["items"], time_lookup)
        matched += sum(1 for item in cat["items"] if "pub_time" in item)
    print(f"Matched pub_time for {matched} articles")

    result = {
        "updated_at": now,
        "categories": data.get("categories", []),
    }

    total = sum(len(c["items"]) for c in result["categories"])
    os.makedirs(os.path.dirname(OUTPUT), exist_ok=True)
    with open(OUTPUT, "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False, indent=2)

    print(f"Written {total} research briefs -> {OUTPUT}")
    for c in result["categories"]:
        print(f"  {c['name']}: {len(c['items'])} items")


if __name__ == "__main__":
    run()