from __future__ import annotations

import argparse
import sys
from datetime import datetime, timezone
from pathlib import Path

import pandas as pd
import yfinance as yf


DEFAULT_EXPIRATION_LIMIT = 10


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Export Yahoo Finance option chains for a symbol to CSV.",
    )
    parser.add_argument("symbol", help="Ticker symbol, e.g., AAPL")
    parser.add_argument(
        "-e",
        "--expiration",
        help="Specific expiration date in YYYY-MM-DD as listed by Yahoo Finance.",
    )
    parser.add_argument(
        "--all-expirations",
        action="store_true",
        help=(
            "Download every listed expiration (default is the first 10 when"
            " no --expiration is provided)."
        ),
    )
    parser.add_argument(
        "--calls-only",
        action="store_true",
        help="Include only call options.",
    )
    parser.add_argument(
        "--puts-only",
        action="store_true",
        help="Include only put options.",
    )
    args = parser.parse_args(argv)

    if args.calls_only and args.puts_only:
        parser.error("Use only one of --calls-only or --puts-only.")

    return args


def pick_expirations(
    ticker: yf.Ticker,
    requested: str | None,
    include_all: bool,
    default_limit: int | None,
) -> tuple[list[str], list[str]]:
    available = list(ticker.options or [])
    if not available:
        raise RuntimeError("No option expirations found for that symbol.")

    if requested:
        if requested not in available:
            raise ValueError(
                f"Expiration {requested!r} not in available dates: {', '.join(available)}"
            )
        targets = [requested]
    elif include_all:
        targets = available
    else:
        limit = default_limit if default_limit is not None else 1
        targets = available[:limit]

    return targets, available


def fetch_option_frames(
    ticker: yf.Ticker, expirations: list[str], include_calls: bool, include_puts: bool
) -> pd.DataFrame:
    frames: list[pd.DataFrame] = []

    for exp in expirations:
        chain = ticker.option_chain(exp)

        if include_calls:
            calls = chain.calls.copy()
            if not calls.empty:
                calls["type"] = "call"
                calls["expiration"] = exp
                frames.append(calls)

        if include_puts:
            puts = chain.puts.copy()
            if not puts.empty:
                puts["type"] = "put"
                puts["expiration"] = exp
                frames.append(puts)

    if not frames:
        return pd.DataFrame()

    combined = pd.concat(frames, ignore_index=True)
    combined.sort_values(["expiration", "type", "strike"], inplace=True)
    combined.reset_index(drop=True, inplace=True)
    return combined


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)

    include_calls = not args.puts_only
    include_puts = not args.calls_only
    include_all_expirations = args.all_expirations
    default_limit = (
        DEFAULT_EXPIRATION_LIMIT if args.expiration is None and not args.all_expirations else None
    )

    ticker = yf.Ticker(args.symbol)

    try:
        target_expirations, available_expirations = pick_expirations(
            ticker, args.expiration, include_all_expirations, default_limit
        )
    except (ValueError, RuntimeError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1

    try:
        options_df = fetch_option_frames(
            ticker, target_expirations, include_calls, include_puts
        )
    except Exception as exc:  # noqa: BLE001
        print(f"error: failed to fetch options data: {exc}", file=sys.stderr)
        return 1

    if options_df.empty:
        print("error: no option data returned.", file=sys.stderr)
        return 1

    timestamp = datetime.now(timezone.utc).isoformat(timespec="seconds")
    output_path = (
        Path.home()
        / "Desktop"
        / f"options_{args.symbol.upper()}_{timestamp}.csv"
    )
    output_path.parent.mkdir(parents=True, exist_ok=True)
    options_df.to_csv(output_path, index=False)

    print(
        f"Wrote {len(options_df)} rows across {len(target_expirations)} expiration(s) "
        f"to {output_path}"
    )
    return 0


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