import os
import logging
import pandas as pd

# Set up logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)

USERNAME = os.environ.get('USERNAME')
PASSWORD = os.environ.get('PASSWORD')
HOST = os.environ.get('HOST')
PORT = os.environ.get('PORT')
DATABASE_NAME = os.environ.get('DATABASE_NAME')

URL = f"mysql+pymysql://{USERNAME}:{PASSWORD}@{HOST}:{PORT}/{DATABASE_NAME}"

# Read data from the database
logging.info('Reading data from NSE Bhav Copy')
nse_df = pd.read_sql_table('nse_bhav_copy', URL, parse_dates=['DATE'])
logging.info('Reading data from BSE Bhav Copy')
bse_df = pd.read_sql_table('bse_bhav_copy', URL, parse_dates=['DATE'])
logging.info('Reading data from Adj NSE Bhav Copy')
adj_nse_df = pd.read_sql_table('adj_nse_bhav_copy', URL, parse_dates=['DATE'])
logging.info('Reading data from Adj BSE Bhav Copy')
adj_bse_df = pd.read_sql_table('adj_bse_bhav_copy', URL, parse_dates=['DATE'])


def get_stock_data(name, start, end):
    """
    Get stock data from the database
    name: str - name of the stock
    start: str - start date for the stock data
    end: str - end date for the stock data
    """
    is_nse = name.endswith('.NS')
    if is_nse:
        name = name[:-3]
        
    adj_df = adj_nse_df if is_nse else adj_bse_df
    df = nse_df if is_nse else bse_df

    df = adj_df if name in adj_df['SYMBOL'].values else df

    df = df.query('SYMBOL == @name')

    if start and end:
        df = df.loc[(df['DATE'] >= start) & (df['DATE'] <= end)]
    elif start:
        df = df.loc[df['DATE'] >= start]
    elif end:
        df = df.loc[df['DATE'] <= end]

    return df