#!/usr/bin/env python3
"""Fetch and map company business-segment revenue exposure.

Source: Eastmoney public BusinessAnalysis endpoint. The data is a structured
vendor representation of company disclosures, not a page-level annual-report
verification. Raw responses are cached per security for audit and restart.
"""

from __future__ import annotations

import argparse
import gzip
import json
import re
import time
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import numpy as np
import pandas as pd


ROOT = Path(__file__).resolve().parents[1]
RAW = ROOT / "data" / "research" / "raw" / "business_segments_eastmoney"
OUT = ROOT / "data" / "research" / "derived"
ENDPOINT = "https://emweb.securities.eastmoney.com/PC_HSF10/BusinessAnalysis/PageAjax"

TARGETS = {
    "房地产开发": r"房地产|商品房|住宅开发|商业地产|地产开发|物业开发",
    "基础设施建设": r"基础设施|基建|工程承包|工程施工|市政工程|铁路工程|公路工程|桥梁工程|轨道交通工程|水利工程",
    "煤炭": r"煤炭|原煤|动力煤|焦煤|洗精煤|煤化工",
    "钢铁": r"钢铁|钢材|板材|线材|不锈钢|特钢|钢管|螺纹钢|热轧|冷轧",
    "软件与信息技术服务": r"软件|信息技术服务|系统集成|云服务|数字化服务|IT服务",
    "通信服务": r"通信服务|电信服务|移动通信|宽带|云网|IDC|数据中心",
    "半导体": r"半导体|集成电路|芯片|晶圆|封装测试|封测|光刻|刻蚀|硅片|功率器件|存储器",
    "新能源汽车": r"新能源汽车|新能源车|电动车|电动汽车|动力电池|电驱|电控|充电桩|车用锂电池",
    "光伏风电储能": r"光伏|太阳能|风电|风力发电|风机|储能|逆变器|光伏组件|电池片|硅料|硅棒",
}
CLASS_LABELS = {"1": "按行业分类", "2": "按产品分类", "3": "按地区分类"}
CLASS_PRIORITY = {"按产品分类": 0, "按行业分类": 1}


def eastmoney_code(thscode: str) -> str:
    code, suffix = thscode.split(".", 1)
    prefix = {"SH": "SH", "SZ": "SZ", "BJ": "BJ"}.get(suffix.upper())
    if not prefix:
        raise ValueError(f"Unsupported security suffix: {thscode}")
    return f"{prefix}{code}"


def raw_path(thscode: str) -> Path:
    return RAW / f"{thscode.replace('.', '_')}.json.gz"


def read_raw(path: Path) -> dict:
    with gzip.open(path, "rt", encoding="utf-8") as handle:
        return json.load(handle)


def write_raw(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with gzip.open(path, "wt", encoding="utf-8") as handle:
        json.dump(payload, handle, ensure_ascii=False, separators=(",", ":"))


def fetch_one(thscode: str, refresh: bool, retries: int = 4) -> tuple[str, str, int]:
    path = raw_path(thscode)
    if path.exists() and not refresh:
        payload = read_raw(path)
        return thscode, "缓存", len(payload.get("zygcfx", []))
    url = ENDPOINT + "?" + urllib.parse.urlencode({"code": eastmoney_code(thscode)})
    request = urllib.request.Request(
        url,
        headers={
            "User-Agent": "Mozilla/5.0",
            "Referer": "https://emweb.securities.eastmoney.com/",
            "Accept-Encoding": "gzip",
        },
    )
    for attempt in range(retries):
        try:
            with urllib.request.urlopen(request, timeout=30) as response:
                body = response.read()
                if response.headers.get("Content-Encoding") == "gzip":
                    body = gzip.decompress(body)
            payload = json.loads(body.decode("utf-8"))
            write_raw(path, payload)
            return thscode, "成功", len(payload.get("zygcfx", []))
        except Exception as exc:
            if attempt == retries - 1:
                return thscode, f"失败:{type(exc).__name__}", 0
            time.sleep(0.5 * (2**attempt))
    raise AssertionError("unreachable")


def parse_payload(thscode: str, payload: dict) -> list[dict]:
    rows: list[dict] = []
    for item in payload.get("zygcfx", []):
        rows.append(
            {
                "证券代码": thscode,
                "报告日期": item.get("REPORT_DATE"),
                "分类类型": CLASS_LABELS.get(str(item.get("MAINOP_TYPE")), str(item.get("MAINOP_TYPE"))),
                "主营构成": item.get("ITEM_NAME"),
                "主营收入": item.get("MAIN_BUSINESS_INCOME"),
                "收入比例": item.get("MBI_RATIO"),
                "主营成本": item.get("MAIN_BUSINESS_COST"),
                "主营利润": item.get("MAIN_BUSINESS_RPOFIT"),
                "毛利率": item.get("GROSS_RPOFIT_RATIO"),
            }
        )
    return rows


def map_exposure(detail: pd.DataFrame, names: pd.Series) -> tuple[pd.DataFrame, pd.DataFrame]:
    if detail.empty:
        return pd.DataFrame(), pd.DataFrame()
    detail["报告日期"] = pd.to_datetime(detail["报告日期"], errors="coerce")
    for column in ("主营收入", "收入比例", "主营成本", "主营利润", "毛利率"):
        detail[column] = pd.to_numeric(detail[column], errors="coerce")
    detail = detail[
        detail["报告日期"].dt.year.between(1990, 2025)
        & detail["分类类型"].isin(CLASS_PRIORITY)
    ].copy()
    detail["证券简称"] = detail["证券代码"].map(names)

    group_keys = ["证券代码", "报告日期", "分类类型"]
    detail["分类主营收入合计"] = detail.groupby(group_keys, dropna=False)["主营收入"].transform(
        "sum", min_count=1
    )
    mapped_frames: list[pd.DataFrame] = []
    matched_parts: list[pd.DataFrame] = []
    for target, pattern in TARGETS.items():
        matched = detail[
            detail["主营构成"].fillna("").str.contains(pattern, regex=True, na=False)
        ].copy()
        if matched.empty:
            continue
        matched["细分产业"] = target
        matched_parts.append(matched)
        aggregated = (
            matched.groupby(group_keys, dropna=False)
            .agg(
                匹配分项数=("主营构成", "size"),
                匹配主营收入=("主营收入", lambda values: values.sum(min_count=1)),
                分类主营收入合计=("分类主营收入合计", "first"),
                收入比例报告值=("收入比例", lambda values: values.sum(min_count=1)),
                匹配主营构成=(
                    "主营构成",
                    lambda values: "；".join(sorted(set(values.dropna().astype(str)))),
                ),
            )
            .reset_index()
        )
        aggregated["细分产业"] = target
        mapped_frames.append(aggregated)
    mapped = pd.concat(mapped_frames, ignore_index=True) if mapped_frames else pd.DataFrame()
    if mapped.empty:
        return mapped, pd.DataFrame()
    mapped["证券简称"] = mapped["证券代码"].map(names)
    mapped["收入比例计算值"] = np.where(
        mapped["分类主营收入合计"].gt(0),
        mapped["匹配主营收入"] / mapped["分类主营收入合计"],
        np.nan,
    )
    reported_valid = mapped["收入比例报告值"].between(0, 1.05)
    calculated_valid = mapped["收入比例计算值"].between(0, 1.05)
    mapped["收入暴露比例"] = np.where(
        reported_valid,
        mapped["收入比例报告值"],
        np.where(calculated_valid, mapped["收入比例计算值"], np.nan),
    )
    mapped["收入比例质量"] = np.select(
        [reported_valid, calculated_valid],
        ["接口报告值", "按分类收入重算"],
        default="无有效比例",
    )
    mapped["暴露证据状态"] = "第三方结构化主营构成已匹配"
    mapped["年报逐页复核"] = "否"
    mapped["数据来源"] = "东方财富公开主营构成接口"
    mapped["证据边界"] = "底层来自上市公司披露的结构化转录；未逐页核对年报原文"
    mapped = mapped[
        [
            "证券代码",
            "证券简称",
            "报告日期",
            "分类类型",
            "细分产业",
            "匹配分项数",
            "匹配主营收入",
            "分类主营收入合计",
            "收入暴露比例",
            "收入比例计算值",
            "收入比例报告值",
            "收入比例质量",
            "匹配主营构成",
            "暴露证据状态",
            "年报逐页复核",
            "数据来源",
            "证据边界",
        ]
    ]
    mapped["分类优先级"] = mapped["分类类型"].map(CLASS_PRIORITY)
    annual = mapped[mapped["报告日期"].dt.month.eq(12) & mapped["报告日期"].dt.day.eq(31)].copy()
    annual = annual.sort_values(
        ["证券代码", "细分产业", "报告日期", "分类优先级"],
        ascending=[True, True, True, True],
    )
    annual = annual.drop_duplicates(["证券代码", "细分产业", "报告日期"], keep="first")
    annual = annual.drop(columns="分类优先级")
    matched_detail = pd.concat(matched_parts, ignore_index=True) if matched_parts else pd.DataFrame()
    return annual, matched_detail


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--limit", type=int)
    parser.add_argument("--refresh", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    RAW.mkdir(parents=True, exist_ok=True)
    company = pd.read_csv(OUT / "a_share_company_snapshot_ifind.csv.gz")
    universe = company[company["year"] == 2025][["code", "name"]].drop_duplicates("code")
    if args.limit:
        universe = universe.head(args.limit)
    codes = universe["code"].astype(str).tolist()
    names = universe.set_index("code")["name"]

    statuses: list[dict] = []
    with ThreadPoolExecutor(max_workers=args.workers) as executor:
        futures = {
            executor.submit(fetch_one, code, args.refresh): code
            for code in codes
        }
        completed = 0
        for future in as_completed(futures):
            code, status, row_count = future.result()
            statuses.append({"证券代码": code, "请求状态": status, "主营构成行数": row_count})
            completed += 1
            if completed % 250 == 0 or completed == len(codes):
                print(f"business exposure: {completed}/{len(codes)}", flush=True)

    status_frame = pd.DataFrame(statuses)
    status_frame.to_csv(
        OUT / "company_business_exposure_fetch_status.csv", index=False, encoding="utf-8-sig"
    )
    parsed: list[dict] = []
    for code in codes:
        path = raw_path(code)
        if path.exists():
            parsed.extend(parse_payload(code, read_raw(path)))
    detail = pd.DataFrame(parsed)
    exposure, matched_detail = map_exposure(detail, names)
    detail.to_csv(
        OUT / "company_business_segment_history.csv.gz", index=False, compression="gzip"
    )
    exposure.to_csv(
        OUT / "company_subindustry_revenue_exposure.csv", index=False, encoding="utf-8-sig"
    )
    matched_detail.to_csv(
        OUT / "company_subindustry_revenue_exposure_detail.csv.gz",
        index=False,
        compression="gzip",
    )

    successful = status_frame["请求状态"].isin(["成功", "缓存"])
    nonempty = status_frame["主营构成行数"].gt(0)
    latest_exposure = (
        exposure[exposure["报告日期"].dt.year.eq(2025)].copy()
        if not exposure.empty
        else exposure
    )
    coverage_rows = []
    for target in TARGETS:
        target_rows = latest_exposure[latest_exposure["细分产业"] == target]
        coverage_rows.append(
            {
                "细分产业": target,
                "2025全A公司数": len(codes),
                "接口成功公司数": int(successful.sum()),
                "有主营构成公司数": int(nonempty.sum()),
                "2025匹配暴露公司数": int(target_rows["证券代码"].nunique()),
                "2025收入比例有效公司数": int(
                    target_rows.loc[target_rows["收入暴露比例"].notna(), "证券代码"].nunique()
                ),
                "证据层级": "第三方结构化主营构成；未逐页年报复核",
                "数据来源": "东方财富公开主营构成接口",
                "更新时间": pd.Timestamp.now().date().isoformat(),
            }
        )
    coverage = pd.DataFrame(coverage_rows)
    coverage.to_csv(
        OUT / "company_subindustry_revenue_exposure_coverage.csv",
        index=False,
        encoding="utf-8-sig",
    )
    if not exposure.empty:
        annual_coverage = (
            exposure.assign(年份=exposure["报告日期"].dt.year)
            .groupby(["年份", "细分产业"], as_index=False)
            .agg(
                匹配暴露公司数=("证券代码", "nunique"),
                收入比例有效公司数=("收入暴露比例", "count"),
                匹配主营收入=("匹配主营收入", "sum"),
            )
        )
        annual_coverage["数据来源"] = "东方财富公开主营构成接口"
        annual_coverage["证据边界"] = "按2025年全A公司池回溯；不含此前已退市且未在2025年股票池的公司"
        annual_coverage.to_csv(
            OUT / "company_subindustry_revenue_exposure_annual_coverage.csv",
            index=False,
            encoding="utf-8-sig",
        )

    assert len(codes) == status_frame["证券代码"].nunique()
    assert set(coverage["细分产业"]) == set(TARGETS)
    assert (coverage["接口成功公司数"] <= coverage["2025全A公司数"]).all()
    if not exposure.empty:
        assert exposure["收入暴露比例"].dropna().ge(0).all()
        assert exposure["收入暴露比例"].dropna().le(1.05).all()
        assert not exposure.duplicated(["证券代码", "细分产业", "报告日期"]).any()
    print("business-exposure assertions: OK")


if __name__ == "__main__":
    main()