chinese-equity-quant/plugins/joinquant/cli.py

"""CLI commands for the JoinQuant comparison plugin."""

from __future__ import annotations

import click

from plugins.joinquant.browser import (
    browser_snapshot,
    load_env_file,
    run_browser_backtest,
    run_browser_sim_trade,
    save_login_state,
    save_login_state_from_env,
    write_browser_config_template,
)
from plugins.joinquant.export_targets import export_targets
from plugins.joinquant.ingest import ingest_joinquant_outputs
from plugins.joinquant.reconcile import reconcile_joinquant
from plugins.joinquant.smoke import prepare_smoke_test
from plugins.joinquant.wrapper_strategy import write_wrapper_strategy


@click.group(name="joinquant")
def joinquant():
    """Compare internal portfolio simulation with JoinQuant output."""


@joinquant.command("export-targets")
@click.option("--positions-path", required=True, help="Portfolio positions parquet from `portfolio build`")
@click.option("--portfolio-name", required=True, help="Portfolio run to export")
@click.option(
    "--mode",
    "mode",
    type=click.Choice(["target_shares", "target_value"]),
    default="target_shares",
    show_default=True,
    help="JoinQuant target order mode",
)
@click.option("--start-date", default=None, help="Inclusive YYYY-MM-DD start date")
@click.option("--end-date", default=None, help="Inclusive YYYY-MM-DD end date")
@click.option(
    "--execution-calendar-path",
    default=None,
    help="Daily data parquet/dataset used to shift position dates to next execution session",
)
@click.option("--out-dir", default="plugins_output/joinquant/targets", show_default=True)
@click.option("--force", is_flag=True, help="Overwrite frozen target/snapshot files")
def export_targets_cmd(
    positions_path,
    portfolio_name,
    mode,
    start_date,
    end_date,
    execution_calendar_path,
    out_dir,
    force,
):
    """Export frozen daily target files for JoinQuant."""
    snapshots = export_targets(
        positions_path=positions_path,
        portfolio_name=portfolio_name,
        mode=mode,
        start_date=start_date,
        end_date=end_date,
        execution_calendar_path=execution_calendar_path,
        out_dir=out_dir,
        force=force,
    )
    click.echo(f"Exported JoinQuant targets: {len(snapshots)} day(s)")
    for snapshot in snapshots:
        click.echo(
            f"  {snapshot['date']}: {snapshot['n_symbols']} symbols, "
            f"sha256={str(snapshot['file_sha256'])[:12]}"
        )


@joinquant.command("ingest")
@click.option("--portfolio-name", required=True, help="Portfolio run name")
@click.option("--fills-csv", required=True, help="JoinQuant fills CSV")
@click.option("--positions-csv", required=True, help="JoinQuant positions CSV")
@click.option("--pnl-csv", required=True, help="JoinQuant daily PnL CSV")
@click.option("--out-dir", default="plugins_output/joinquant/ingested", show_default=True)
def ingest_cmd(portfolio_name, fills_csv, positions_csv, pnl_csv, out_dir):
    """Normalize JoinQuant CSV exports to parquet."""
    paths = ingest_joinquant_outputs(
        portfolio_name=portfolio_name,
        fills_csv=fills_csv,
        positions_csv=positions_csv,
        pnl_csv=pnl_csv,
        out_dir=out_dir,
    )
    click.echo(f"Saved JoinQuant fills:     {paths['fills']}")
    click.echo(f"Saved JoinQuant positions: {paths['positions']}")
    click.echo(f"Saved JoinQuant pnl:       {paths['pnl']}")


@joinquant.command("reconcile")
@click.option("--portfolio-name", required=True, help="Portfolio run name")
@click.option("--targets-dir", required=True, help="Directory containing exported daily target files")
@click.option("--our-fills-path", required=True, help="Internal simulator fills parquet")
@click.option("--our-positions-path", required=True, help="Internal portfolio positions parquet")
@click.option("--our-pnl-path", required=True, help="Internal simulator PnL parquet")
@click.option("--jq-fills-path", required=True, help="Normalized JoinQuant fills parquet")
@click.option("--jq-positions-path", required=True, help="Normalized JoinQuant positions parquet")
@click.option("--jq-pnl-path", required=True, help="Normalized JoinQuant PnL parquet")
@click.option("--out-dir", default="plugins_output/joinquant/reconcile", show_default=True)
@click.option("--share-tolerance", default=0.0, show_default=True, type=float)
@click.option("--price-rel-tolerance", default=1e-4, show_default=True, type=float)
@click.option("--pnl-tolerance", default=1.0, show_default=True, type=float)
@click.option("--booksize", default=None, type=float, help="Booksize for value tolerance inference")
def reconcile_cmd(
    portfolio_name,
    targets_dir,
    our_fills_path,
    our_positions_path,
    our_pnl_path,
    jq_fills_path,
    jq_positions_path,
    jq_pnl_path,
    out_dir,
    share_tolerance,
    price_rel_tolerance,
    pnl_tolerance,
    booksize,
):
    """Write per-symbol and daily JoinQuant reconciliation reports."""
    paths = reconcile_joinquant(
        portfolio_name=portfolio_name,
        targets_dir=targets_dir,
        our_fills_path=our_fills_path,
        our_positions_path=our_positions_path,
        our_pnl_path=our_pnl_path,
        jq_fills_path=jq_fills_path,
        jq_positions_path=jq_positions_path,
        jq_pnl_path=jq_pnl_path,
        out_dir=out_dir,
        share_tolerance=share_tolerance,
        price_rel_tolerance=price_rel_tolerance,
        pnl_tolerance=pnl_tolerance,
        booksize=booksize,
    )
    click.echo(f"Saved reconciliation parquet: {paths['daily_reconcile']}")
    click.echo(f"Saved reconciliation summary: {paths['summary_md']}")
    click.echo(f"Saved reconciliation CSV:     {paths['summary_csv']}")


@joinquant.command("write-wrapper")
@click.option("--portfolio-name", required=True, help="Portfolio run name")
@click.option(
    "--mode",
    "mode",
    type=click.Choice(["target_shares", "target_value"]),
    default="target_shares",
    show_default=True,
)
@click.option("--out-path", required=True, help="Path for generated standalone strategy")
@click.option("--allow-short", is_flag=True, help="Do not clip negative targets in the generated wrapper")
@click.option(
    "--embed-targets-dir",
    default=None,
    help="Embed CSV target files from this directory into the strategy source",
)
def write_wrapper_cmd(portfolio_name, mode, out_path, allow_short, embed_targets_dir):
    """Generate a standalone JoinQuant wrapper strategy."""
    path = write_wrapper_strategy(
        portfolio_name=portfolio_name,
        mode=mode,
        out_path=out_path,
        allow_short=allow_short,
        embedded_targets_dir=embed_targets_dir,
    )
    click.echo(f"Saved JoinQuant wrapper strategy: {path}")


@joinquant.command("prepare-smoke")
@click.option("--out-dir", required=True, help="Root directory for generated smoke-test artifacts")
@click.option(
    "--universe",
    default="sh600000,sz000001,sh600519,sz002594,sz300750",
    show_default=True,
    help="Universe for the tiny real-data download",
)
@click.option("--trade-symbol", default="sh600000", show_default=True)
@click.option("--start-date", default="2024-01-02", show_default=True)
@click.option("--end-date", default="2024-01-12", show_default=True)
@click.option("--portfolio-name", default="jq_smoke_one_stock_long", show_default=True)
@click.option("--shares", default=1000, show_default=True, type=int)
@click.option("--booksize", default=1_000_000.0, show_default=True, type=float)
@click.option("--max-signal-dates", default=3, show_default=True, type=int)
@click.option("--cost-bps", default=5.0, show_default=True, type=float)
@click.option("--slippage-bps", default=5.0, show_default=True, type=float)
@click.option("--volume-frac", default=0.02, show_default=True, type=float)
@click.option("--force", is_flag=True, help="Overwrite existing frozen target files")
def prepare_smoke_cmd(
    out_dir,
    universe,
    trade_symbol,
    start_date,
    end_date,
    portfolio_name,
    shares,
    booksize,
    max_signal_dates,
    cost_bps,
    slippage_bps,
    volume_frac,
    force,
):
    """Prepare a one-command local real-data JoinQuant smoke test."""
    manifest = prepare_smoke_test(
        out_dir=out_dir,
        universe=universe,
        trade_symbol=trade_symbol,
        start_date=start_date,
        end_date=end_date,
        portfolio_name=portfolio_name,
        shares=shares,
        booksize=booksize,
        max_signal_dates=max_signal_dates,
        cost_bps=cost_bps,
        slippage_bps=slippage_bps,
        volume_frac=volume_frac,
        force=force,
    )
    click.echo(f"Prepared JoinQuant smoke manifest: {manifest['manifest_path']}")
    click.echo(f"Wrapper: {manifest['wrapper_path']}")
    click.echo(f"Targets: {manifest['targets_dir']}")
    click.echo(f"Expected JoinQuant exports: {manifest['joinquant_export_dir']}")
    summary = manifest["local_summary"]
    click.echo(
        f"Local simulator: {summary['n_pnl_rows']} days, "
        f"PnL={summary['total_pnl']:,.2f}, cost={summary['total_cost']:,.2f}, "
        f"blocked={summary['blocked_trades']}"
    )


@joinquant.command("browser-login")
@click.option(
    "--storage-state",
    default="~/.config/chinese-equity-quant/joinquant_storage_state.json",
    show_default=True,
    help="Path where the authenticated browser state is stored",
)
@click.option(
    "--login-url",
    default="https://www.joinquant.com/user/login/index",
    show_default=True,
)
@click.option("--headless", is_flag=True, help="Use a headless browser")
@click.option(
    "--wait-seconds",
    default=0,
    show_default=True,
    type=int,
    help="Seconds to wait before saving state; 0 prompts for Enter after login",
)
def browser_login_cmd(storage_state, login_url, headless, wait_seconds):
    """Open JoinQuant login and save reusable browser session state."""
    path = save_login_state(
        storage_state=storage_state,
        login_url=login_url,
        headless=headless,
        wait_seconds=wait_seconds,
    )
    click.echo(f"Saved JoinQuant browser state: {path}")


@joinquant.command("browser-login-env")
@click.option(
    "--env-path",
    default="~/.config/chinese-equity-quant/joinquant.env",
    show_default=True,
    help="Env file with JOINQUANT_USERNAME and JOINQUANT_PASSWORD",
)
@click.option(
    "--storage-state",
    default="~/.config/chinese-equity-quant/joinquant_storage_state.json",
    show_default=True,
    help="Path where the authenticated browser state is stored",
)
@click.option(
    "--login-url",
    default="https://www.joinquant.com/user/login/index",
    show_default=True,
)
@click.option("--headed", is_flag=True, help="Run with a visible browser")
@click.option("--out-dir", default=None, help="Login artifact directory")
@click.option("--timeout-ms", default=120_000, show_default=True, type=int)
def browser_login_env_cmd(env_path, storage_state, login_url, headed, out_dir, timeout_ms):
    """Try credential-based JoinQuant login from an env file."""
    env = load_env_file(env_path)
    missing = [
        key for key in ["JOINQUANT_USERNAME", "JOINQUANT_PASSWORD"]
        if not env.get(key)
    ]
    if missing:
        raise click.ClickException(f"Missing required env keys: {', '.join(missing)}")
    report = save_login_state_from_env(
        env_path=env_path,
        storage_state=storage_state,
        login_url=login_url,
        headless=not headed,
        out_dir=out_dir,
        timeout_ms=timeout_ms,
    )
    click.echo(f"Saved login report: {report['report_path']}")
    if report["logged_in"]:
        click.echo(f"Saved JoinQuant browser state: {report['storage_state']}")
    else:
        raise click.ClickException(
            "JoinQuant login did not complete. Check the saved screenshot/HTML; "
            "CAPTCHA/SMS/2FA may require `joinquant browser-login` once."
        )


@joinquant.command("write-browser-config")
@click.option("--out-path", required=True, help="Path for JSON browser automation config")
@click.option("--strategy-url", default="", help="JoinQuant strategy edit/backtest/模拟盘 URL")
@click.option(
    "--flow",
    type=click.Choice(["backtest", "sim-trade"]),
    default="backtest",
    show_default=True,
    help="Browser automation template to write",
)
def write_browser_config_cmd(out_path, strategy_url, flow):
    """Write a browser automation selector/action config template."""
    path = write_browser_config_template(out_path, strategy_url=strategy_url, flow=flow)
    click.echo(f"Saved JoinQuant browser config template: {path}")


@joinquant.command("browser-snapshot")
@click.option("--url", required=True, help="Logged-in JoinQuant page URL to inspect")
@click.option(
    "--storage-state",
    default="~/.config/chinese-equity-quant/joinquant_storage_state.json",
    show_default=True,
)
@click.option("--out-dir", required=True, help="Directory for page.html and page.png")
@click.option("--headed", is_flag=True, help="Run with a visible browser")
@click.option("--timeout-ms", default=60_000, show_default=True, type=int)
def browser_snapshot_cmd(url, storage_state, out_dir, headed, timeout_ms):
    """Save a page screenshot and HTML for selector discovery."""
    paths = browser_snapshot(
        url=url,
        storage_state=storage_state,
        out_dir=out_dir,
        headless=not headed,
        timeout_ms=timeout_ms,
    )
    click.echo(f"Saved HTML:       {paths['html']}")
    click.echo(f"Saved screenshot: {paths['screenshot']}")


@joinquant.command("run-browser-backtest")
@click.option("--manifest-path", required=True, help="Manifest from `joinquant prepare-smoke`")
@click.option("--config-path", required=True, help="JSON browser automation config")
@click.option(
    "--storage-state",
    default="~/.config/chinese-equity-quant/joinquant_storage_state.json",
    show_default=True,
)
@click.option("--out-dir", default=None, help="Browser-run artifact directory")
@click.option("--headed", is_flag=True, help="Run with a visible browser")
@click.option("--no-auto-reconcile", is_flag=True, help="Skip ingest/reconcile after downloads")
def run_browser_backtest_cmd(
    manifest_path,
    config_path,
    storage_state,
    out_dir,
    headed,
    no_auto_reconcile,
):
    """Run JoinQuant browser automation from manifest and config."""
    report = run_browser_backtest(
        manifest_path=manifest_path,
        config_path=config_path,
        storage_state=storage_state,
        out_dir=out_dir,
        headless=not headed,
        auto_reconcile=not no_auto_reconcile,
    )
    click.echo(f"Saved browser run report: {report['report_path']}")
    if report["downloaded"]:
        click.echo("Downloaded JoinQuant CSVs:")
        for key, path in report["downloaded"].items():
            click.echo(f"  {key}: {path}")
    if report["reconcile_paths"]:
        click.echo(f"Saved reconciliation summary: {report['reconcile_paths']['summary_md']}")


@joinquant.command("run-browser-sim")
@click.option("--manifest-path", required=True, help="Manifest from target preparation")
@click.option("--config-path", required=True, help="JSON simulated-trading browser config")
@click.option(
    "--storage-state",
    default="~/.config/chinese-equity-quant/joinquant_storage_state.json",
    show_default=True,
)
@click.option("--out-dir", default=None, help="Browser-run artifact directory")
@click.option("--headed", is_flag=True, help="Run with a visible browser")
@click.option("--no-auto-reconcile", is_flag=True, help="Skip ingest/reconcile after downloads")
def run_browser_sim_cmd(
    manifest_path,
    config_path,
    storage_state,
    out_dir,
    headed,
    no_auto_reconcile,
):
    """Run JoinQuant 模拟盘 browser automation from manifest and config."""
    report = run_browser_sim_trade(
        manifest_path=manifest_path,
        config_path=config_path,
        storage_state=storage_state,
        out_dir=out_dir,
        headless=not headed,
        auto_reconcile=not no_auto_reconcile,
    )
    click.echo(f"Saved browser sim-trade report: {report['report_path']}")
    if report["downloaded"]:
        click.echo("Downloaded JoinQuant CSVs:")
        for key, path in report["downloaded"].items():
            click.echo(f"  {key}: {path}")
    if report["reconcile_paths"]:
        click.echo(f"Saved reconciliation summary: {report['reconcile_paths']['summary_md']}")