#!/usr/bin/env python3
"""
专利采集器 - 从 Google Patents 采集专利结构化数据
支持按标题搜索和按专利号直接采集
依赖: scrapling >= 0.4.2
"""

import argparse
import json
import os
import re
import sys
import time
import urllib.parse
from pathlib import Path

# ─────────────────────── 采集字段定义 ───────────────────────

PATENT_FIELDS = [
    "中文名称",
    "英文名称",
    "摘要",
    "专利类型",
    "申请/专利号",
    "专利日期",
    "公开/公告号",
    "公开/公告日",
    "主分类号",
    "分类号",
    "发明/设计人",
    "优先权",
    "法律状态",
    "专利说明",
    "专利授权日期",
    "所属国家",
]

# ─────────────────────── Google Patents 搜索 ───────────────────────

def build_search_url(query, page=0):
    """构建 Google Patents 搜索 URL"""
    encoded = urllib.parse.quote(query)
    return f"https://patents.google.com/?q={encoded}&page={page}"


def build_patent_url(patent_id, lang="en"):
    """构建 Google Patents 专利详情 URL"""
    return f"https://patents.google.com/patent/{patent_id}/{lang}"


# ─────────────────────── 专利详情解析 ───────────────────────

def parse_patent_detail(md_text, patent_url=""):
    """从 Google Patents Markdown 页面解析专利结构化数据"""
    record = {f: "" for f in PATENT_FIELDS}
    record["来源URL"] = patent_url

    # ---- 英文名称（两种模式） ----
    # 模式1: 页面顶部 "PATENT_ID - Title - Google Patents"
    title_match = re.search(
        r"^[A-Z]{2}\d+[A-Z]\d?\s*-\s*(.+?)\s*-\s*Google Patents",
        md_text, re.M,
    )
    if title_match:
        record["英文名称"] = title_match.group(1).strip()
    else:
        # 模式2: h2 标题（跳过导航性短标题）
        for m in re.finditer(r"^(.+?)\n={3,}\s*$", md_text, re.M):
            t = m.group(1).strip()
            if len(t) > 10 and t.lower() not in (
                "patents", "abstract", "claims", "classifications",
            ):
                record["英文名称"] = t
                break

    # ---- 摘要 ----
    abs_match = re.search(
        r"### Abstract translated from\s*\n+(.*?)(?=\n###|\n##|\nImages|\nClassifications|\nClaims|\Z)",
        md_text,
        re.S,
    )
    if abs_match:
        record["摘要"] = abs_match.group(1).strip().replace("\n", " ")
    else:
        # fallback: "Abstract" 段落
        abs_match2 = re.search(
            r"\bAbstract\b\s*\n+(.*?)(?=\n###|\n##|\nImages|\nClassifications|\nClaims|\Z)",
            md_text,
            re.S,
        )
        if abs_match2:
            record["摘要"] = abs_match2.group(1).strip().replace("\n", " ")

    # ---- 发明人 ----
    inventors = []
    inv_section = re.search(r"Inventor\s*\n((?:\s*:\s*\[.*?\]\(#\)\s*\n)+)", md_text)
    if inv_section:
        inventors = re.findall(r"\[\s*(.*?)\s*\]\(#\)", inv_section.group(1))
    record["发明/设计人"] = "; ".join(inventors)

    # ---- 当前受让人（Assignee） ----
    assignee_match = re.search(
        r"Current Assignee.*?\n\s*:\s*(.+?)(?:\n|$)", md_text
    )
    if assignee_match:
        record["专利说明"] = f"受让人: {assignee_match.group(1).strip()}"

    # ---- 申请号 / 公开号 / 日期 / 法律状态 ----
    # 提取所有 worldwide applications 块
    app_pattern = re.compile(
        r"Application number:\s*(.+?)\s*\n\s*Filing date:\s*(\d{4}-\d{2}-\d{2})\s*\n\s*Legal status:\s*(.+?)(?:\n|$)",
        re.M,
    )
    apps = []
    for m in app_pattern.finditer(md_text):
        apps.append({
            "app_number": m.group(1).strip(),
            "filing_date": m.group(2).strip(),
            "legal_status": m.group(3).strip(),
        })

    # US application 优先
    us_apps = [a for a in apps if "US" in a["app_number"].upper()]
    primary_app = us_apps[0] if us_apps else (apps[0] if apps else {})

    if primary_app:
        record["申请/专利号"] = primary_app["app_number"]
        record["专利日期"] = primary_app["filing_date"]
        record["法律状态"] = primary_app["legal_status"]

    # ---- 优先权 (Priority date) ----
    # 尝试多种模式提取优先权日期
    # 模式1: 在 Worldwide applications 之前可能有 Priority date 行
    # 模式2: 最早的 Filing date 即为优先权日
    priority_match = re.search(
        r"Priority date\s*\n(\d{4}-\d{2}-\d{2})", md_text
    )
    if priority_match:
        record["优先权"] = priority_match.group(1)
    elif apps:
        # 最早的 filing date 作为优先权日
        earliest = min(a["filing_date"] for a in apps)
        record["优先权"] = earliest

    # ---- 公开/公告号 & 公开/公告日 ----
    # 从 URL 或页面标题中提取专利号
    url_match = re.search(r"/patent/([A-Z]{2}\d+[A-Z]\d?)", patent_url)
    if url_match:
        record["公开/公告号"] = url_match.group(1)
    else:
        # 从页面首行提取
        pub_match = re.match(r"([A-Z]{2}\d+[A-Z]\d?)", md_text[:200])
        if pub_match:
            record["公开/公告号"] = pub_match.group(1)

    # Publication date - 尝试多种模式
    # 模式1: "YYYY-MM-DD\n[Publication of PATENT_ID]"（最可靠）
    pid_for_match = record.get("公开/公告号", "")
    if pid_for_match:
        pub_date_match = re.search(
            r"(\d{4}-\d{2}-\d{2})\s*\n\s*\[Publication of " + re.escape(pid_for_match) + r"\]",
            md_text
        )
        if pub_date_match:
            record["公开/公告日"] = pub_date_match.group(1)
    # 模式2: "Application granted\nYYYY-MM-DD\n[Publication of...]"
    if not record["公开/公告日"]:
        grant_pub_match = re.search(
            r"Application granted\s*\n\s*(\d{4}-\d{2}-\d{2})",
            md_text
        )
        if grant_pub_match:
            record["公开/公告日"] = grant_pub_match.group(1)
    # 模式3: "Publication date" 表头格式
    if not record["公开/公告日"]:
        pub_date_match = re.search(
            r"Publication [Dd]ate[:\s]*\n\s*(\d{4}-\d{2}-\d{2})", md_text
        )
        if pub_date_match:
            record["公开/公告日"] = pub_date_match.group(1)

    # ---- 分类号 (CPC / IPC) ----
    cpc_codes = re.findall(r"\[([A-Z]\d{2}[A-Z]\s*\d+/\d+)\]\(#\)", md_text)
    if cpc_codes:
        record["分类号"] = "; ".join(dict.fromkeys(cpc_codes))  # 去重保序
        record["主分类号"] = cpc_codes[0]

    # ---- 专利类型 ----
    patent_id = record.get("公开/公告号", "") or record.get("申请/专利号", "")
    if patent_id:
        if "B1" in patent_id or "B2" in patent_id:
            record["专利类型"] = "授权专利"
        elif "A1" in patent_id:
            record["专利类型"] = "发明专利申请公开"
        elif "A2" in patent_id:
            record["专利类型"] = "专利申请"
        else:
            record["专利类型"] = "专利"

    # ---- 专利授权日期 ----
    grant_date_match = re.search(
        r"(?:Grant date|Date of Patent|Granted)[:\s]*\n\s*(\d{4}-\d{2}-\d{2})", md_text
    )
    if grant_date_match:
        record["专利授权日期"] = grant_date_match.group(1)
    elif record["专利类型"] == "授权专利" and record["公开/公告日"]:
        record["专利授权日期"] = record["公开/公告日"]
    # 如果是授权专利(B1/B2)但无日期，尝试从美国数据中提取
    if not record["专利授权日期"] and record["专利类型"] == "授权专利":
        # 如果有公开/公告日，授权日期通常相同或接近
        if record["公开/公告日"]:
            record["专利授权日期"] = record["公开/公告日"]
        else:
            # 尝试提取 "Application granted" 后的日期
            grant_match = re.search(
                r"Application granted\s*\n\s*(\d{4}-\d{2}-\d{2})", md_text
            )
            if grant_match:
                record["专利授权日期"] = grant_match.group(1)

    # ---- 所属国家 ----
    if patent_id:
        country_code = re.match(r"([A-Z]{2})", patent_id)
        if country_code:
            code = country_code.group(1)
            country_map = {
                "US": "美国", "EP": "欧洲", "WO": "世界知识产权组织",
                "GB": "英国", "DE": "德国", "FR": "法国", "JP": "日本",
                "CN": "中国", "KR": "韩国", "CA": "加拿大", "AU": "澳大利亚",
            }
            record["所属国家"] = country_map.get(code, code)
    if not record["所属国家"] and us_apps:
        record["所属国家"] = "美国"

    return record


# ─────────────────────── 搜索结果解析 ───────────────────────

def parse_search_results(md_text):
    """从搜索结果页提取专利号列表"""
    patent_ids = []
    # 匹配专利号格式: US10620304B2, EP1234567A1, WO2020123456A1 等
    pattern = re.compile(r"\[([A-Z]{2}\d+[A-Z]\d?)\]\(https://patentimages")
    for m in pattern.finditer(md_text):
        pid = m.group(1)
        if pid not in patent_ids:
            patent_ids.append(pid)
    return patent_ids


# ─────────────────────── PDF 下载 ───────────────────────

def build_pdf_url(patent_id):
    """构建专利 PDF 下载 URL（Google Patents 格式）"""
    # US10620304B2 → US10620304.pdf
    # 格式: https://patentimages.storage.googleapis.com/XX/XX/XX/XXXXX/USXXXXXXX.pdf
    # 实际需要从详情页提取 PDF URL
    return None  # 需要从详情页中提取


def extract_pdf_url(md_text):
    """从专利详情页 Markdown 中提取 PDF 链接"""
    pdf_match = re.search(
        r"https://patentimages\.storage\.googleapis\.com/[a-f0-9/]+/[A-Z]{2}\d+\.pdf",
        md_text,
    )
    return pdf_match.group(0) if pdf_match else None


# ─────────────────────── 主采集函数 ───────────────────────

def collect_patent_by_title(title, scrapling_bin="scrapling", output_dir="/tmp/patents"):
    """按标题搜索并采集专利"""
    os.makedirs(output_dir, exist_ok=True)
    results = []

    # Step 1: 搜索
    search_url = build_search_url(f'TI=("{title}")')
    search_file = os.path.join(output_dir, f"search_{hash(title) % 100000}.md")

    print(f"[搜索] {title}")
    os.system(
        f'{scrapling_bin} extract stealthy-fetch "{search_url}" "{search_file}" '
        f"--network-idle --wait 5000 --timeout 60000"
    )

    if not os.path.exists(search_file):
        print(f"  [失败] 搜索结果未保存")
        return results

    with open(search_file, "r", encoding="utf-8") as f:
        search_md = f.read()

    patent_ids = parse_search_results(search_md)
    print(f"  找到 {len(patent_ids)} 个专利号: {patent_ids[:5]}")

    # Step 2: 逐个采集详情
    for pid in patent_ids:
        record = collect_patent_detail(pid, scrapling_bin, output_dir)
        if record:
            results.append(record)
        time.sleep(2)  # 礼貌延迟

    # 清理临时搜索文件
    if os.path.exists(search_file):
        os.remove(search_file)

    return results


def collect_patent_detail(patent_id, scrapling_bin="scrapling", output_dir="/tmp/patents"):
    """采集单个专利详情页"""
    patent_url = build_patent_url(patent_id)
    detail_file = os.path.join(output_dir, f"detail_{patent_id}.md")

    print(f"  [采集] {patent_id} ...")
    ret = os.system(
        f'{scrapling_bin} extract stealthy-fetch "{patent_url}" "{detail_file}" '
        f"--network-idle --wait 3000 --timeout 60000"
    )

    if ret != 0 or not os.path.exists(detail_file):
        print(f"    [失败] 详情页采集失败")
        return None

    with open(detail_file, "r", encoding="utf-8") as f:
        detail_md = f.read()

    record = parse_patent_detail(detail_md, patent_url)
    if not record.get("英文名称"):
        print(f"    [跳过] 无法解析专利标题")
        return None

    # 提取 PDF URL 并下载
    pdf_url = extract_pdf_url(detail_md)
    if pdf_url:
        record["PDF_URL"] = pdf_url
        pdf_dir = os.path.join(output_dir, "pdfs")
        os.makedirs(pdf_dir, exist_ok=True)
        pdf_file = os.path.join(pdf_dir, f"{patent_id}.pdf")
        print(f"    [下载PDF] {pdf_url}")
        os.system(f'curl -sL "{pdf_url}" -o "{pdf_file}"')
        if os.path.exists(pdf_file) and os.path.getsize(pdf_file) > 1000:
            record["PDF本地路径"] = pdf_file
            print(f"    [完成] PDF 已保存: {pdf_file}")
        else:
            print(f"    [警告] PDF 下载可能失败")

    # 清理临时详情文件
    if os.path.exists(detail_file):
        os.remove(detail_file)

    return record


# ─────────────────────── 批量采集 ───────────────────────

def batch_collect(patent_list_file, output_dir="/tmp/patents", scrapling_bin="scrapling"):
    """从列表文件批量采集"""
    with open(patent_list_file, "r", encoding="utf-8") as f:
        lines = [l.strip() for l in f if l.strip() and not l.startswith("#")]

    all_results = []
    for i, line in enumerate(lines, 1):
        print(f"\n{'='*60}")
        print(f"[{i}/{len(lines)}] {line}")
        print(f"{'='*60}")

        # 判断是专利号还是标题搜索
        if re.match(r"^[A-Z]{2}\d+", line):
            # 直接专利号
            record = collect_patent_detail(line, scrapling_bin, output_dir)
            if record:
                all_results.append(record)
        else:
            # 按标题搜索
            results = collect_patent_by_title(line, scrapling_bin, output_dir)
            all_results.extend(results)

        time.sleep(3)  # 礼貌延迟

    return all_results


# ─────────────────────── 数据护照 ───────────────────────

def generate_data_passport(records, query_info=""):
    """生成数据护照"""
    passport = {
        "合规性声明": "所有数据均从 Google Patents 公开页面采集，遵守 robots.txt 协议",
        "采集时间": time.strftime("%Y-%m-%d %H:%M:%S %Z"),
        "数据来源": "Google Patents (patents.google.com)",
        "采集工具": f"开源情报-专利采集器 / Scrapling",
        "采集范围": f"共采集 {len(records)} 条专利记录",
        "查询信息": query_info,
        "数据质量评分": "待评估",
        "字段覆盖度": {},
    }

    # 计算字段覆盖度
    for field in PATENT_FIELDS:
        filled = sum(1 for r in records if r.get(field))
        passport["字段覆盖度"][field] = f"{filled}/{len(records)} ({filled/max(len(records),1)*100:.0f}%)"

    return passport


# ─────────────────────── 输出 ───────────────────────

def save_results(records, output_dir, query_info=""):
    """保存采集结果"""
    os.makedirs(output_dir, exist_ok=True)

    # JSON 输出
    json_file = os.path.join(output_dir, "patents.json")
    with open(json_file, "w", encoding="utf-8") as f:
        json.dump(records, f, ensure_ascii=False, indent=2)
    print(f"\n[保存] JSON: {json_file}")

    # CSV 输出
    csv_file = os.path.join(output_dir, "patents.csv")
    import csv
    with open(csv_file, "w", newline="", encoding="utf-8-sig") as f:
        writer = csv.DictWriter(f, fieldnames=PATENT_FIELDS + ["来源URL", "PDF_URL", "PDF本地路径"])
        writer.writeheader()
        for r in records:
            writer.writerow({k: r.get(k, "") for k in PATENT_FIELDS + ["来源URL", "PDF_URL", "PDF本地路径"]})
    print(f"[保存] CSV: {csv_file}")

    # 数据护照
    passport = generate_data_passport(records, query_info)
    passport_file = os.path.join(output_dir, "data_passport.json")
    with open(passport_file, "w", encoding="utf-8") as f:
        json.dump(passport, f, ensure_ascii=False, indent=2)
    print(f"[保存] 数据护照: {passport_file}")

    return json_file, csv_file, passport_file


# ─────────────────────── CLI 入口 ───────────────────────

def main():
    parser = argparse.ArgumentParser(description="专利采集器 - 从 Google Patents 采集结构化数据")
    parser.add_argument("--title", type=str, help="按标题搜索专利")
    parser.add_argument("--patent-id", type=str, help="直接采集指定专利号")
    parser.add_argument("--list", type=str, help="从文件批量采集（每行一个标题或专利号）")
    parser.add_argument("--output", type=str, default="/tmp/patents", help="输出目录")
    parser.add_argument("--scrapling", type=str, default="scrapling", help="scrapling 可执行文件路径")

    args = parser.parse_args()

    if not any([args.title, args.patent_id, args.list]):
        parser.print_help()
        sys.exit(1)

    if args.patent_id:
        records = []
        record = collect_patent_detail(args.patent_id, args.scrapling, args.output)
        if record:
            records.append(record)
        save_results(records, args.output, f"专利号: {args.patent_id}")

    elif args.title:
        records = collect_patent_by_title(args.title, args.scrapling, args.output)
        save_results(records, args.output, f"标题搜索: {args.title}")

    elif args.list:
        records = batch_collect(args.list, args.output, args.scrapling)
        save_results(records, args.output, f"批量采集: {args.list}")

    print(f"\n[完成] 共采集 {len(records)} 条专利记录")


if __name__ == "__main__":
    main()