#!/usr/bin/env python3
"""Small raw-data adapter for the AKShare endpoints used by Fund Lab."""

import argparse
import datetime as dt
import json
import re
import sys
from decimal import Decimal, InvalidOperation

import akshare as ak
import pandas as pd


SCHEMA_VERSION = "fund-lab.akshare.v1"


def collected_at():
    return dt.datetime.now(dt.timezone.utc).isoformat().replace("+00:00", "Z")


def source_revision():
    return f"akshare-{getattr(ak, '__version__', 'unknown')}/eastmoney"


def is_missing(value):
    try:
        return bool(pd.isna(value))
    except (TypeError, ValueError):
        return False


def text(value):
    if is_missing(value):
        return None
    value = str(value).strip()
    return value or None


def fund_code(value):
    value = text(value)
    if value is None:
        return None
    if value.isdigit() and len(value) < 6:
        value = value.zfill(6)
    return value if re.fullmatch(r"\d{6}", value) else None


def decimal_text(value):
    value = text(value)
    if value is None:
        return None
    try:
        return format(Decimal(value), "f")
    except InvalidOperation:
        return None


def date_text(value):
    if is_missing(value):
        return None
    if isinstance(value, (dt.datetime, dt.date)):
        return value.strftime("%Y-%m-%d")
    parsed = pd.to_datetime(value, errors="coerce")
    if pd.isna(parsed):
        return None
    return parsed.strftime("%Y-%m-%d")


def search(query):
    query = query.strip()
    if not query:
        raise ValueError("search query cannot be empty")

    frame = ak.fund_name_em()
    query_lower = query.casefold()
    rows = []

    for _, row in frame.iterrows():
        code = fund_code(row.get("基金代码"))
        name = text(row.get("基金简称"))
        pinyin = text(row.get("拼音缩写"))
        full_pinyin = text(row.get("拼音全称"))

        if code is None or name is None:
            continue

        searchable = [code, name, pinyin or "", full_pinyin or ""]
        if not any(query_lower in value.casefold() for value in searchable):
            continue

        rows.append(
            {
                "code": code,
                "name": name,
                "fund_type": text(row.get("基金类型")),
                "rank": 0 if code == query else (1 if name == query else 2),
            }
        )

    rows.sort(key=lambda item: (item.pop("rank"), item["code"]))
    return {
        "schema_version": SCHEMA_VERSION,
        "operation": "search",
        "source": "akshare",
        "source_revision": source_revision(),
        "collected_at": collected_at(),
        "instruments": rows[:100],
    }


def nav(code):
    if not re.fullmatch(r"\d{6}", code):
        raise ValueError("fund code must contain exactly six digits")

    frame = ak.fund_open_fund_info_em(
        symbol=code,
        indicator="单位净值走势",
        period="成立来",
    )
    observations = []

    for _, row in frame.iterrows():
        nav_date = date_text(row.get("净值日期"))
        nav_value = decimal_text(row.get("单位净值"))
        if nav_date is None or nav_value is None:
            continue

        observations.append(
            {
                "nav_date": nav_date,
                "published_at": None,
                "nav": nav_value,
                "accumulated_nav": None,
                "daily_return": decimal_text(row.get("日增长率")),
            }
        )

    if not observations:
        raise ValueError(f"AKShare returned no usable NAV observations for {code}")

    return {
        "schema_version": SCHEMA_VERSION,
        "operation": "nav",
        "source": "akshare",
        "source_revision": source_revision(),
        "collected_at": collected_at(),
        "instrument": {"code": code},
        "observations": observations,
    }


def bond_maturity_date(code):
    # Maturity is a best-effort enrichment: the quote endpoint has no expiry
    # date, so a failure here only means maturity_date stays null.
    try:
        frame = ak.bond_zh_cov_info(symbol=code, indicator="基本信息")
        if frame is None or frame.empty:
            return None
        return date_text(frame.iloc[0].get("EXPIRE_DATE"))
    except Exception:
        return None


def bond_quote_from_spot(code):
    # Realtime convertible-bond snapshot (sina). Gives name/trade but no date or
    # accrued-interest fields; returns None when the code carries no valid price.
    frame = ak.bond_zh_hs_cov_spot()
    if frame is None or frame.empty:
        return None

    prefix_rank = {"sh": 0, "sz": 1, "bj": 2}
    candidates = []

    for _, row in frame.iterrows():
        symbol = text(row.get("symbol")) or ""
        row_code = text(row.get("code")) or ""
        if row_code != code and not symbol.endswith(code):
            continue

        price = decimal_text(row.get("trade"))
        if price is None or Decimal(price) <= 0:
            continue

        prefix = symbol[:2].lower()
        candidates.append((prefix_rank.get(prefix, 3), symbol, text(row.get("name")), price))

    if not candidates:
        return None

    _, symbol, name, price = sorted(candidates, key=lambda item: (item[0], item[1]))[0]
    revision = source_revision()

    return {
        "source": "akshare",
        "source_revision": f"{revision}/hs-cov-spot",
        "name": name,
        "price": price,
        "clean_price": price,
        "accrued_interest": None,
        "date": None,
        "maturity_date": bond_maturity_date(code),
        "symbol": symbol,
    }


def bond_quote_from_daily(code):
    # Exchange bond daily history (sina). The venue is not derivable from the
    # six-digit code alone, so try both prefixes; the first that returns rows
    # wins. Gives a settled close and the trade date but no accrued-interest
    # breakdown.
    for prefix in ("sh", "sz"):
        symbol = f"{prefix}{code}"

        try:
            frame = ak.bond_zh_hs_daily(symbol=symbol)
        except Exception:
            continue

        if frame is None or frame.empty or "date" not in frame.columns or "close" not in frame.columns:
            continue

        rows = []

        for _, row in frame.iterrows():
            row_date = date_text(row.get("date"))
            close = decimal_text(row.get("close"))
            if row_date is None or close is None or Decimal(close) <= 0:
                continue
            rows.append((row_date, close))

        if not rows:
            continue

        rows.sort(key=lambda item: item[0])
        row_date, close = rows[-1]
        revision = source_revision()

        return {
            "source": "akshare",
            "source_revision": f"{revision}/hs-bond-daily",
            "name": None,
            "price": close,
            "clean_price": close,
            "accrued_interest": None,
            "date": row_date,
            "maturity_date": bond_maturity_date(code),
            "symbol": symbol,
        }

    return None


def bond_quote(code):
    code = text(code)
    if code is None or not re.fullmatch(r"\d{6}", code):
        raise ValueError("bond code must contain exactly six digits")

    for attempt in (bond_quote_from_spot, bond_quote_from_daily):
        try:
            quote = attempt(code)
        except Exception:
            quote = None

        if quote is not None:
            return {
                "schema_version": SCHEMA_VERSION,
                "operation": "bond-quote",
                "source": quote["source"],
                "source_revision": quote["source_revision"],
                "collected_at": collected_at(),
                "instrument": {"code": code},
                "quote": {
                    "name": quote["name"],
                    "price": quote["price"],
                    "clean_price": quote["clean_price"],
                    "accrued_interest": quote["accrued_interest"],
                    "date": quote["date"],
                    "maturity_date": quote["maturity_date"],
                },
            }

    raise ValueError(f"no valid bond quote for {code}")


def first_positive_decimal(*values):
    for value in values:
        price = decimal_text(value)
        if price is not None and Decimal(price) > 0:
            return price
    return None


def stock_price_from_bid_ask(code):
    # Single-symbol realtime snapshot (eastmoney). Returns None when the venue
    # rejects the symbol or carries no valid last price.
    frame = ak.stock_bid_ask_em(symbol=code)
    if frame is None or frame.empty:
        return None

    for _, row in frame.iterrows():
        item = text(row.get("item"))
        if item in ("最新", "最新价"):
            return first_positive_decimal(row.get("value"))

    return None


def stock_from_snapshot(code):
    # Whole-market snapshots differ per provider; each entry is
    # (fetch, code column, name column, price column, exchange prefixes).
    providers = [
        (ak.stock_zh_a_spot_tx, "code", "name", "zxj", {"sh": 0, "sz": 1, "bj": 2}),
        (ak.stock_zh_a_spot_em, "代码", "名称", "最新价", None),
    ]

    for fetch, code_key, name_key, price_key, prefixes in providers:
        try:
            frame = fetch()
        except Exception:
            continue

        if frame is None or frame.empty:
            continue

        candidates = []

        for _, row in frame.iterrows():
            row_code = text(row.get(code_key)) or ""
            if row_code != code and not row_code.lower().endswith(code):
                continue

            price = first_positive_decimal(row.get(price_key))
            if price is None:
                continue

            prefix = row_code[:2].lower()
            rank = prefixes.get(prefix, 3) if prefixes is not None else 0
            candidates.append((rank, row_code, text(row.get(name_key)), price))

        if candidates:
            _, _, name, price = sorted(candidates, key=lambda item: (item[0], item[1]))[0]
            return name, price

    return None


def stock_daily_observations(frame, date_key, close_key, volume_key, amount_key):
    rows = []

    for _, row in frame.iterrows():
        row_date = date_text(row.get(date_key))
        close = decimal_text(row.get(close_key))
        if row_date is None or close is None or Decimal(close) <= 0:
            continue

        rows.append(
            {
                "date": row_date,
                "close": close,
                "volume": decimal_text(row.get(volume_key)),
                "amount": decimal_text(row.get(amount_key)),
            }
        )

    rows.sort(key=lambda item: item["date"])
    return rows


def stock_market_symbol(code):
    if code[0] in ("6", "9"):
        return f"sh{code}"
    if code[0] in ("0", "3", "2"):
        return f"sz{code}"
    if code[0] in ("4", "8"):
        return f"bj{code}"
    return f"sh{code}"


def stock_daily(code, days_text):
    code = text(code)
    if code is None or not re.fullmatch(r"\d{6}", code):
        raise ValueError("stock code must contain exactly six digits")

    if days_text is None:
        days = 5
    else:
        try:
            days = int(days_text)
        except (TypeError, ValueError):
            raise ValueError("days must be an integer between 1 and 30")

    if days < 1 or days > 30:
        raise ValueError("days must be an integer between 1 and 30")

    symbol = stock_market_symbol(code)
    providers = [
        ((ak.stock_zh_a_hist, {"symbol": code, "period": "daily", "adjust": ""}), "日期", "收盘", "成交量", "成交额"),
        ((ak.stock_zh_a_daily, {"symbol": symbol, "adjust": ""}), "date", "close", "volume", "amount"),
        ((ak.stock_zh_a_hist_tx, {"symbol": symbol}), "date", "close", "volume", "amount"),
    ]

    for (fetch, kwargs), date_key, close_key, volume_key, amount_key in providers:
        try:
            frame = fetch(**kwargs)
        except Exception:
            continue

        if frame is None or frame.empty:
            continue

        observations = stock_daily_observations(frame, date_key, close_key, volume_key, amount_key)
        if not observations:
            continue

        return {
            "schema_version": SCHEMA_VERSION,
            "operation": "stock-daily",
            "source": "akshare",
            "source_revision": source_revision(),
            "collected_at": collected_at(),
            "instrument": {"code": code},
            "observations": observations[-days:],
        }

    raise ValueError(f"no valid daily history for {code}")


def stock_quote(code):
    code = text(code)
    if code is None or not re.fullmatch(r"\d{6}", code):
        raise ValueError("stock code must contain exactly six digits")

    name = None
    price = None

    try:
        price = stock_price_from_bid_ask(code)
    except Exception:
        price = None

    if price is None:
        found = stock_from_snapshot(code)
        if found is None:
            raise ValueError(f"no valid quote for {code}")
        name, price = found

    return {
        "schema_version": SCHEMA_VERSION,
        "operation": "stock-quote",
        "source": "akshare",
        "source_revision": source_revision(),
        "collected_at": collected_at(),
        "instrument": {"code": code},
        "quote": {
            "name": name,
            "price": price,
            "currency": "CNY",
        },
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--operation", choices=("search", "nav", "bond-quote", "stock-quote", "stock-daily"), required=True)
    parser.add_argument("--query")
    parser.add_argument("--code")
    parser.add_argument("--days")
    args = parser.parse_args()

    try:
        if args.operation == "search":
            payload = search(args.query)
        elif args.operation == "nav":
            payload = nav(args.code)
        elif args.operation == "bond-quote":
            payload = bond_quote(args.code)
        elif args.operation == "stock-quote":
            payload = stock_quote(args.code)
        else:
            payload = stock_daily(args.code, args.days)

        json.dump(payload, sys.stdout, ensure_ascii=False, separators=(",", ":"))
        sys.stdout.write("\n")
        return 0
    except Exception as error:
        print(f"AKShare collector failed: {error}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())