#!/usr/bin/env python

from sys import exit
from dotenv import load_dotenv
from typing import NoReturn
from paperone.utils import (
    parse_date_yyyymmdd,
    days_range_from,
    format_date_readable,
    is_trading_day,
)
from os import environ
from paperone.taapi import TaapiClient, IndicatorEnum
from rich.progress import (
    Progress,
    SpinnerColumn,
    TextColumn,
    BarColumn,
    TaskProgressColumn,
    TimeRemainingColumn,
)

load_dotenv()


def main() -> NoReturn:
    api_key = environ.get("API_KEY")

    if not api_key:
        print("API_KEY not set")
        exit(0)

    date = parse_date_yyyymmdd("20250821")
    days_range = 60
    dates_range = days_range_from(date, days_range)
    tickers = ["VIX"]
    indicators = list(IndicatorEnum)

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TaskProgressColumn(),
        TimeRemainingColumn(),
    ) as progress:
        # Overall ticker progress
        ticker_task = progress.add_task(
            "[cyan]Processing tickers...", total=len(tickers)
        )

        with TaapiClient(api_key) as client:
            for ticker in [x for x in tickers if x in client.get_available_tickers()]:
                # Update ticker task
                progress.update(
                    ticker_task, description=f"[cyan]Processing {ticker}..."
                )

                # Price loading subtask
                price_task = progress.add_task(
                    f"[green]  └─ Loading prices for {ticker}...",
                    total=len(dates_range),
                )

                prices = {}
                for d in dates_range:
                    progress.update(
                        price_task,
                        description=f"[green]  └─ Loading {ticker} price for {format_date_readable(d)}...",
                        advance=1,
                    )
                    result = client.query_price_on_day(ticker, d)
                    if result:
                        prices[d.day] = result.value

                # Remove price task when done
                progress.remove_task(price_task)

                # Indicator loading subtask
                indicator_task = progress.add_task(
                    f"[yellow]  └─ Loading indicators for {ticker}...",
                    total=len(indicators),
                )

                for indicator_enum in indicators:
                    progress.update(
                        indicator_task,
                        description=f"[yellow]  └─ Loading {ticker} indicator: {indicator_enum.name}...",
                        advance=1,
                    )

                    try:
                        indicator_results = client.query_indicator(
                            ticker, indicator_enum.value, date, results=days_range
                        )
                    except Exception as e:
                        progress.console.print(
                            f"[red]Error retrieving {indicator_enum.name}: {e}"
                        )
                        continue

                    if not indicator_results:
                        continue

                    progress.console.print(
                        f"\n[bold]{ticker} - {indicator_enum.name}:[/bold]"
                    )

                    trading_day_values = [
                        x for x in indicator_results if is_trading_day(x.datetime)
                    ]

                    for r in trading_day_values:
                        price_str = (
                            f"${prices[r.datetime.day]:.2f}"
                            if r.datetime.day in prices
                            else "N/A"
                        )
                        progress.console.print(
                            f"  {format_date_readable(r.datetime)} ({price_str}) - {indicator_enum.name}: {r.value:.2f}"
                        )

                # Remove indicator task when done
                progress.remove_task(indicator_task)

                # Advance overall ticker progress
                progress.advance(ticker_task)

    exit(0)


if __name__ == "__main__":
    main()