import random
import numpy as np
from typing import Literal
from datetime import datetime, timedelta

import pandas as pd
from hamilton.htypes import Parallelizable, Collect
from hamilton.function_modifiers import source, parameterized_subdag, parameterize_sources


def _create_dummy_trades_data(num_trades: int = 100, exchange: Literal['NYSE', 'NASDAQ', 'AMEX'] = 'NYSE'):
    # Define the columns for the trades dataframe
    columns = ['trade_id', 'timestamp', 'symbol', 'quantity', 'price', 'trade_type', 'exchange']
    
    # Generate dummy data
    trade_ids = range(1, num_trades + 1)
    start_date = datetime.now() - timedelta(days=30)
    timestamps = [start_date + timedelta(seconds=random.randint(0, 2592000)) for _ in range(num_trades)]  # Random timestamps within the last 30 days
    symbols = ['AAPL', 'GOOGL', 'MSFT', 'AMZN', 'TSLA']
    symbol_choices = [random.choice(symbols) for _ in range(num_trades)]
    quantities = np.random.randint(1, 1000, size=num_trades)  # Random quantities between 1 and 1000
    prices = np.round(np.random.uniform(100, 1500, size=num_trades), 2)  # Random prices between 100 and 1500
    trade_types = ['buy', 'sell']
    trade_type_choices = [random.choice(trade_types) for _ in range(num_trades)]
    exchange_choices = [exchange for _ in range(num_trades)]
    
    # Create the dataframe
    data = {
        'trade_id': trade_ids,
        'timestamp': timestamps,
        'symbol': symbol_choices,
        'quantity': quantities,
        'price': prices,
        'trade_type': trade_type_choices,
        'exchange': exchange_choices
    }
    
    trades_df = pd.DataFrame(data, columns=columns)
    
    return trades_df


def file_path(file_paths: list) -> Parallelizable[str]:
    for f in file_paths:
        yield f

def raw_data_src1(file_path: str) -> pd.DataFrame:
    return pd.read_csv(file_path)

def hist_data(raw_data_src1: Collect[pd.DataFrame]) -> pd.DataFrame:
    return pd.concat(raw_data_src1)


@parameterized_subdag(
    file_path, raw_data_src1, hist_data,
    from_datasource_1={"inputs" : {"file_paths" : source("INPUT_FILES_DATASOURCE_1")}},
    from_datasource_2={"inputs" : {"file_paths" : source("INPUT_FILES_DATASOURCE_2")}},
    from_datasource_3={"inputs" : {"file_paths" : source("INPUT_FILES_DATASOURCE_3")}}
)
def etl_pipeline(hist_data: pd.DataFrame) -> pd.DataFrame:
    return hist_data


def norm_data(
    from_datasource_1: pd.DataFrame, 
    from_datasource_2: pd.DataFrame, 
    from_datasource_3: pd.DataFrame
) -> pd.DataFrame:
    return pd.concat(locals().values(), axis="index")



if __name__ == "__main__":
    import os
    from glob import glob
    from hamilton import driver, base
    import __main__

    # Example usage
    # create dummy data
    def create_dummy_files() -> bool:
        exchanges = ('NYSE', 'NASDAQ', 'AMEX')
        # dont create dummy files if exist
        print("dummy files exist")
        if any(f.startswith(exchanges) for f in os.listdir()):
            return True

        print("creating dummy files")
        for exch in exchanges:
            print(exch)
            for _ in range(2):
                _create_dummy_trades_data(1000, exchange=exch).to_csv(f"{exch}_{datetime.now():%Y%m%d%H%M%S%f}.csv",index=False)
        return False
    
    create_dummy_files()

    dr = (
        driver.Builder()
        .enable_dynamic_execution(allow_experimental_mode=True)
        .with_modules(__main__)
        .with_adapters(base.PandasDataFrameResult())
        .build()
    )

    print("executing dag")
    result = dr.execute(
        [   
            "norm_data",
        ]
        ,inputs=dict(
            INPUT_FILES_DATASOURCE_1=glob('NYSE*.csv'),
            INPUT_FILES_DATASOURCE_2=glob('NASDAQ*.csv'),
            INPUT_FILES_DATASOURCE_3=glob('AMEX*.csv'),
        )
    )
    
    print(result)