'''
This script gets daily as well as new data for the stocks in the database "bse_data1" -> OHCL from Zerodha Demat API
'''

from z_getBSEdata import getBSEdata

# System Libs
import datetime
import json

# Installed Libraries
from pymongo import MongoClient
import mysql.connector as mc
from time import sleep
import pandas as pd
import requests

def get_credentials(env='live', filepath=r'/var/www/html/trade_iq/db_updater/credentials.json'):
    """
    Access credentials from a JSON file.

    Args:
    - env (str): Environment to fetch the credentials from ('local' or 'live'). Default is 'local'.
    - filepath (str): Path to the JSON file containing credentials. Default is 'creds.json'.

    Returns:
    - dict: Credentials for the specified environment.
    """
    try:
        # Open and load the JSON file
        with open(filepath, 'r') as f:
            data = json.load(f)

        # Fetch credentials based on environment
        credentials = data.get(env)

        if credentials:
            return credentials
        else:
            raise ValueError(f"No credentials found for environment: {env}")

    except FileNotFoundError:
        print(f"File not found: {filepath}")
    except json.JSONDecodeError:
        print(f"Error decoding JSON in file: {filepath}")
    except Exception as e:
        print(f"An error occurred: {e}")

def prev_close_updaterBSE():
    creds = get_credentials()

    # Credentials
    mongoDB_creds = creds.get("mongoDB_creds")
    connection_url = mongoDB_creds.get("connection_str")
    mongo_database = mongoDB_creds.get("database")
    bse_collection = mongoDB_creds.get("bse_collection")

    client = MongoClient(connection_url)
    db = client[mongo_database]
    collection = db[bse_collection]

    # Fetch all documents
    documents = collection.find()

    # Loop through each document
    for doc in documents:
        previous_close = None

        # Convert the date keys to datetime objects and sort them in ascending order
        sorted_dates = sorted(doc['prices'].keys(), key=lambda date: datetime.datetime.strptime(date, '%Y-%m-%d'))

        # Flag to track if any updates are made
        update_needed = False

        # Iterate through the sorted dates
        for date in sorted_dates:
            data = doc['prices'][date]

            # Check if 'prev_close' is already present
            if 'prev_close' not in data:
                if previous_close is not None:
                    # Add the prev_close field with the previous date's close value
                    data['prev_close'] = previous_close
                    update_needed = True  # Set the flag to True since an update is made
                    print(f"Updated prev_close for date {date}: {previous_close}")

            # Store the current date's close value to use in the next iteration
            previous_close = data.get('close')

        # Update the document in MongoDB only if there were any updates
        if update_needed:
#            collection.update_one({'_id': doc['_id']}, {'$set': {'prices': doc['prices']}})
            pass

    print("prev_close fields updated where missing!")

def bse_fetcher():
    creds = get_credentials()

    # Credentials
    mongoDB_creds = creds.get("mongoDB_creds")
    connection_url = mongoDB_creds.get("connection_str")
    mongo_database = mongoDB_creds.get("database")
    bse_collection = mongoDB_creds.get("bse_collection")

    sqlDB_creds = creds.get("sqlDB_creds")
    user = sqlDB_creds.get("user")
    host = sqlDB_creds.get("host")
    password = sqlDB_creds.get("password")
    sql_database = sqlDB_creds.get("database")
    # port = sqlDB_creds.get("port")

    # Connect to MongoDB
    mongo_client = MongoClient(connection_url)
    db = mongo_client[mongo_database]
    collection = db[bse_collection]

    # Connect to MySQL
    conn = mc.connect(user=user, host=host, password=password, database=sql_database)
    cursor = conn.cursor()

    # Read symbols from MySQL
    cursor.execute("SELECT `Symbol.1` FROM mytable WHERE `exchange` = 'BSE';")
    checkSymbols = cursor.fetchall()

    today = datetime.date.today()
    creds_collection = db["zerodha_credentials"]
    get_creds = creds_collection.find_one({ "date": str(today) })
    AUTH = get_creds["enctoken"]
    print(AUTH)

    # Set up headers for API request
    headers = {
    'accept': '*/*',
    'accept-language': 'en-US,en;q=0.9',
    'authorization': f'enctoken {AUTH}'
    }

    today = datetime.datetime.today().strftime("%Y-%m-%d")
    today = datetime.datetime.strptime(today, "%Y-%m-%d")
    yesterday = (today - datetime.timedelta(days=1)).date()
    default_start_date = datetime.date(2020, 1, 20)  # Set default start date to January 1, 2014

    for symbol in checkSymbols:
        # Check if the symbol exists in MongoDB
        existing_document = collection.find_one({'symbol': symbol[0]})

        if existing_document:
            # If exists, find the last recorded date
            try:
                last_date = max(datetime.datetime.strptime(date, '%Y-%m-%d').date() for date in existing_document['prices'].keys())
                from_date = last_date + datetime.timedelta(days=1)
            except ValueError:
                from_date = default_start_date

            if from_date > yesterday:
                print(f"Data for {symbol[0]} is up to date.")
                continue
        else:
            # If not exists, set from_date to the default start date
            from_date = default_start_date

        print(f"Processing {symbol[0]} from {from_date} to {yesterday}")

        # Get instrument token from MySQL
        # cursor.execute(f"SELECT instrument_token FROM `mytable` WHERE `Symbol.1` LIKE '{symbol[0]}' AND `exchange` LIKE 'BSE' ORDER BY `instrument_type` DESC")
        cursor.execute(f"SELECT instrument_token, segment, exchange_token FROM `mytable` WHERE `Symbol.1` LIKE '{symbol[0]}' AND `exchange` LIKE 'BSE';")
        token = cursor.fetchone()

        if token is not None:
            # instToken = token[0]
            instToken, segment, exchange_token = token

            url = f"https://kite.zerodha.com/oms/instruments/historical/{instToken}/day?user_id=GMG829&oi=0&from={from_date}&to={yesterday}"
#            url = f"https://kite.zerodha.com/oms/instruments/historical/{instToken}/day?user_id=GMG829&oi=0&from=2019-01-01&to=2023-12-31"
            response = json.loads(requests.request("GET", url, headers=headers, data={}).text)
            print(response)

            try:
                if response and 'data' in response and 'candles' in response['data']:
                    new_prices = {}
                    for candle in response['data']['candles']:
                        candle_date = datetime.datetime.strptime(candle[0].split('T')[0], '%Y-%m-%d').date()
                        ohlc = {
                            'open': candle[1],
                            'high': candle[2],
                            'low': candle[3],
                            'close': candle[4],
                            'volume': candle[5]
                        }
                        new_prices[candle_date.strftime('%Y-%m-%d')] = ohlc

                    if existing_document:
                        # Update existing document
#                        existing_document['prices'].update(new_prices)
                        existing_document['prices'].update(new_prices)
                        existing_document['timestamp'] = datetime.datetime.now()
                        collection.replace_one({'symbol': symbol[0]}, existing_document)
                    else:
                        # Insert new document
                        new_document = {
                            'symbol': symbol[0],
                            'Instrument': segment,
                            'exchange_token': int(exchange_token),
                            'instrument_token': int(instToken),
                            'timestamp': datetime.datetime.now(),
                            'exchange': 'BSE',
                            'prices': new_prices
                        }
                        collection.insert_one(new_document)

                    print(f"Completed processing {symbol[0]}")
                else:
                    print(f"No new data available for {symbol[0]}.")
            except Exception as e:
                print(f"Error processing {symbol[0]}: {e}")
        else:
            print(f"No instrument token found for {symbol[0]}")

    print("Processing completed. Results saved to MongoDB")
    mongo_client.close()
    conn.close()

if __name__ == "__main__":
    bse_fetcher()
    sleep(3)
    prev_close_updaterBSE()
    sleep(3)
    getBSEdata()

