import datetime
from pydantic import BaseModel, Field
from typing import Dict, List, Optional
import yfinance as yf
import plotly.express as px
from workcell.integrations.types import PlotlyPlot


class Input(BaseModel):
    tickers: List[str] = Field(
        default=['AAPL','AMZN','META'], max_items=10, description="List of ticker values"
    )


def load_data(tickers):
    """Download ticker price data from ticker list.
    e.g. tickers = ['AAPL','AMZN','GOOG']
    """
    start = datetime.datetime(2022, 1, 1)
    end = datetime.datetime.now() # latest
    data = yf.download(tickers, start=start, end=end, interval='1d')
    # adjust close
    close = data['Adj Close']
    return close


def visualization(df):
    """Visualization price plot by yfinance dataframe.
    e.g. df = yf.download(***)
    """
    template = 'plotly_white'
    fig = px.line(df, x=df.index, y=df.columns.tolist(), template=template)
    fig.update_xaxes(
        rangeslider_visible=True,
        rangeselector=dict(
            buttons=list([
                dict(count=1, label="1m", step="month", stepmode="backward"),
                dict(count=6, label="6m", step="month", stepmode="backward"),
                dict(count=1, label="YTD", step="year", stepmode="todate"),
                dict(count=1, label="1y", step="year", stepmode="backward"),
                dict(step="all")
            ])
        )    
    )
    fig.update_layout(hovermode="x")    
    return fig


def stock_viewer(input: Input) -> PlotlyPlot:
    """Input ticker list, returns multiple stocks price. Data frome yfinance."""
    # Step1. load data
    dataframe = load_data(input.tickers)
    # Step2. create plot
    fig = visualization(dataframe)
    # Step3. wrapped by output
    output = PlotlyPlot(data=fig)
    return output