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

from getNSEdata import getNSEdata
import pandas as pd
import datetime
import requests
import json
from pymongo import MongoClient
import mysql.connector as mc
import os

# Connect to MongoDB
mongo_client = MongoClient('mongodb://192.168.31.176:27017/')
# db = mongo_client['zerodha_data']
db = mongo_client['zerodha_test']
collection = db['nse_data1']

# Connect to MySQL
conn = mc.connect(user='root', host='localhost', password='', database='instruments', port=3307)
cursor = conn.cursor()

# Read the SQLDB
cursor.execute("SELECT `tradingsymbol` FROM zerodha_nse WHERE `exchange` = 'NSE' AND `segment` = 'NSE' AND `instrument_type` = 'EQ';")
checkSymbols = cursor.fetchall()

# Get Auth
AUTH = "NnCItBXGvYU+//QxPIvU4Hy4n4cSaFl/xDD9q6DkJS+7QJIzkSDtCRuiQ1uZphb7x6swDKH9hl3ksVfeOstQTB6pMLSROGXvPYR9PyzcjdRRL5k8kcsBFQ=="
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)
yesterday = (today - datetime.timedelta(days=1)).date()

default_start_date = datetime.date(2014, 1, 1)
# today = datetime.date(2024, 10, 16)  # Set to the current date

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 `tradingsymbol` LIKE '{symbol[0]}' AND `instrument_type` LIKE 'EQ' AND `exchange` LIKE 'NSE' ORDER BY `instrument_type` DESC")
    token = cursor.fetchone()

    if token is not None:
        instToken = token[0]
        print(instToken)

        url = f"https://kite.zerodha.com/oms/instruments/historical/{instToken}/day?user_id=GMG829&oi=0&from={from_date}&to={yesterday}"
        response = json.loads(requests.request("GET", url, headers=headers, data={}).text)
        print(response)  # Check the 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]
                    }
                    new_prices[candle_date.strftime('%Y-%m-%d')] = ohlc

                if existing_document:
                    # Update existing document
                    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],
                        'timestamp': datetime.datetime.now(),
                        '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()