from pathlib import Path

from marketreturns.retrieve import (
    compute_annual_returns,
    compute_monthly_returns,
    download_raw_prices,
    validate_coverage,
)


def main():
    data_dir = Path(__file__).parent / "data"
    data_dir.mkdir(parents=True, exist_ok=True)

    all_data = download_raw_prices(data_dir)

    if not all_data:
        print("No data downloaded. Exiting.")
        return

    validate_coverage(all_data)

    monthly_returns = compute_monthly_returns(all_data)
    annual_returns = compute_annual_returns(all_data)

    returns_dir = data_dir / "returns"
    returns_dir.mkdir(parents=True, exist_ok=True)

    monthly_path = returns_dir / "monthly_returns.csv"
    annual_path = returns_dir / "annual_returns.csv"

    monthly_returns.to_csv(monthly_path, index=False)
    print(f"Monthly returns saved to {monthly_path} ({len(monthly_returns)} rows)")

    annual_returns.to_csv(annual_path, index=False)
    print(f"Annual returns saved to {annual_path} ({len(annual_returns)} rows)")

    print("\nSample annual returns (latest 5 years):")
    print(annual_returns.groupby("index").tail(5).to_string(index=False))


if __name__ == "__main__":
    main()