import json
import websocket
import threading
import time
from datetime import datetime, timedelta
import pickle
import os

FINNHUB_API_KEY = "******************"
STOP_HOUR = 13  # 1 PM in 24-hour format
DATA_FILE = "data_spy_3hrs_250618.pkl"

# Store data
my_data = []
last_timestamp = None  # Keep track of the last timestamp to prevent duplicates

# Load previous data if available
if os.path.exists(DATA_FILE):
    with open(DATA_FILE, "rb") as file:
        my_data = pickle.load(file)
    if my_data:
        last_timestamp = my_data[-1]["timestamp"]  # Get the last recorded timestamp
    print(f"Loaded {len(my_data)} records from previous session.")

# Function to save data
def save_data():
    with open(DATA_FILE, "wb") as file:
        pickle.dump(my_data, file)
    print(f"Saved {len(my_data)} records to {DATA_FILE}")

def on_message(ws, message):
    global my_data, last_timestamp
    data = json.loads(message)

    print(data)

    # Stop collecting data after 1 PM
    if datetime.now().hour >= STOP_HOUR:
        print("1 PM reached, stopping data collection.")
        ws.close()
        return

    if "data" in data:
        for trade in data["data"]:
            trade_timestamp = trade["t"] / 1000  # Convert from ms to seconds

            # Only save new data (avoid duplicates)
            if last_timestamp is None or trade_timestamp > last_timestamp:
                my_data.append({
                    "timestamp": trade_timestamp,
                    "price": trade["p"],
                    "volume": trade["v"]
                })
                last_timestamp = trade_timestamp  # Update last timestamp

def on_error(ws, error):
    print("WebSocket error:", error)

def on_close(ws, close_status_code, close_msg):
    print("WebSocket closed.")
    save_data()  # Save data before closing

    # Only restart the WebSocket if it's before 1 PM
    if datetime.now().hour < STOP_HOUR:
        print("Reconnecting in 1 second...")
        time.sleep(1)
        start_websocket()
    else:
        print("1 PM or later – WebSocket will not restart.")

def on_open(ws):
    ws.send(json.dumps({"type": "subscribe", "symbol": "SPY"}))
    print("Subscribed to SPY")

def start_websocket():
    while datetime.now().hour < STOP_HOUR:
        try:
            ws = websocket.WebSocketApp(
                f"wss://ws.finnhub.io?token={FINNHUB_API_KEY}",
                on_message=on_message,
                on_error=on_error,
                on_close=on_close
            )
            ws.on_open = on_open
            ws.run_forever()
        except Exception as e:
            print(f"Error: {e}. Reconnecting in 1 second...")
            time.sleep(1)

# Wait until 6:30 AM before starting
now = datetime.now()
# target_time = now.replace(hour=6, minute=30, second=0, microsecond=0)
# if now > target_time:
#     target_time += timedelta(days=1)

# print("Waiting until:", target_time)
# time.sleep((target_time - now).total_seconds())

print("Starting WebSocket...")
start_websocket()