import json
import time
import random
import string
from datetime import datetime, timedelta, timezone

# Assume 'client' is an instance of your WebSocket client library
# Example:
# class WebSocketClient:
#     def connect(self): ...
#     def send(self, event, params): ... # Formats and sends the message
#     def recv(self): ... # Receives, parses, and returns message dictionary
#     def close(self): ...
#
# client = WebSocketClient()
# client.connect()

# --- Configuration ---
AUTH_TOKEN = "YOUR_AUTH_TOKEN" # Replace with your valid token
SYMBOL = "OANDA:XAUUSD"
TIMEFRAME = "1" # 1 minute
CHUNK_BARS = 1500 # How many bars to request per chunk (adjust as needed)
TOTAL_CHUNKS = 3 # How many chunks back in time to scrape for this POC

# --- Generate Session IDs (Client-side placeholders) ---
def generate_session_id(prefix):
    return prefix + "_" + "".join(random.choice(string.ascii_letters + string.digits) for _ in range(12))

chart_session_id = generate_session_id("cs")
replay_session_id = generate_session_id("rs")
replay_series_id = generate_session_id("sds_sym_replay") # ID for the resolved replay symbol
series_id_for_replay_data = "s2" # The series ID used after modify_series in logs

# --- Helper Functions ---
def send_message(event, params):
    print(f"--> Sending: {event} {params}")
    client.send(event, params)
    time.sleep(0.1) # Small delay between sends

def receive_message(timeout=5):
    """Receives a message, returns parsed dict or None on timeout/error."""
    try:
        message_str = client.recv(timeout=timeout) # Assuming recv returns raw string or None
        if not message_str:
            print("<-- No message received (timeout)")
            return None
        # Handle potential TradingView heartbeat messages (~h~)
        if message_str.startswith("~h~"):
            print(f"<-- Received Heartbeat: {message_str}")
            # Respond to heartbeat if required by library/protocol
            # client.send_raw(message_str.replace("~h~", "~h~"))
            return {"type": "heartbeat"} # Or None, depending on handling

        # Basic parsing assuming format ~m~<len>~m~{json}
        parts = message_str.split("~m~")
        if len(parts) >= 3 and parts[0] == '' and parts[1].isdigit():
            message_data = json.loads(parts[2])
            print(f"<-- Received: {message_data.get('m', 'UNKNOWN')}")
            return message_data
        else:
            print(f"<-- Received Unknown format: {message_str}")
            return None
    except Exception as e:
        print(f"<-- Error receiving/parsing message: {e}")
        return None

# --- Main Script ---
all_candles = []
current_replay_end_timestamp = int(datetime.now(timezone.utc).timestamp()) # Start from now

try:
    # 1. Initial Setup
    send_message("set_auth_token", [AUTH_TOKEN])
    send_message("chart_create_session", [chart_session_id, ""])
    send_message("switch_timezone", [chart_session_id, "Etc/UTC"])
    send_message("replay_create_session", [replay_session_id])

    # Wait for replay session confirmation (optional but good practice)
    # In a real scenario, you'd wait for 'replay_instance_id' response
    print("Waiting briefly for session setup...")
    time.sleep(2)

    # Resolve symbol ONCE for replay context
    symbol_param = f"={json.dumps({'replay': replay_session_id, 'symbol': {'adjustment': 'splits', 'currency-id': 'USD', 'session': 'regular', 'symbol': SYMBOL}})}"
    send_message("resolve_symbol", [chart_session_id, replay_series_id, symbol_param])

    # Modify series ONCE to display replay data
    # Note: Assuming sds_1 was the initial series created implicitly or explicitly.
    # If not, you might need to create_series first. Using 'sds_1' as per log context.
    send_message("modify_series", [chart_session_id, "sds_1", series_id_for_replay_data, replay_series_id, TIMEFRAME, ""])

    print("Waiting briefly for symbol/series setup...")
    time.sleep(2) # Allow time for server processing

    # 2. Replay Loop for Multiple Chunks
    for i in range(TOTAL_CHUNKS):
        print(f"\n--- Starting Chunk {i+1} ---")
        print(f"Requesting data ending at: {datetime.fromtimestamp(current_replay_end_timestamp, timezone.utc)}")

        # Reset replay to the end of the desired chunk
        replay_point_id = f"rp_{i+1}"
        send_message("replay_reset", [replay_session_id, replay_point_id, current_replay_end_timestamp])

        # Replay Add Series - *Potentially only needed once* after replay_create_session
        # If the server keeps the series association across resets, skip this in loop
        # For robustness based on logs, let's include it but use a consistent ID
        replay_symbol_req_id = f"rsr_{i+1}"
        symbol_definition = f"={json.dumps({'adjustment': 'splits', 'currency-id': 'USD', 'session': 'regular', 'symbol': SYMBOL})}"
        send_message("replay_add_series", [replay_session_id, replay_symbol_req_id, symbol_definition, TIMEFRAME])

        # Request data backwards
        send_message("request_more_data", [chart_session_id, series_id_for_replay_data, CHUNK_BARS])

        # Receive data for this chunk
        chunk_candles = []
        oldest_timestamp_in_chunk = current_replay_end_timestamp
        hit_limit = False
        receive_start_time = time.time()

        while time.time() - receive_start_time < 30: # 30-second timeout per chunk
            msg = receive_message()
            if msg is None: continue # Timeout or error
            if msg.get("type") == "heartbeat": continue

            # Check for candle data
            if msg.get("m") == "timescale_update":
                # Payload structure: msg['p'][1][series_id]['s']
                if (len(msg.get("p", [])) > 1 and
                        isinstance(msg["p"][1], dict) and
                        series_id_for_replay_data in msg["p"][1]):
                    series_data = msg["p"][1][series_id_for_replay_data]
                    if "s" in series_data and isinstance(series_data["s"], list):
                        candles = series_data["s"]
                        print(f"Received {len(candles)} candles for chunk {i+1}")
                        for candle in candles:
                            # Structure: {"i": index, "v": [ts, o, h, l, c, vol]}
                            if isinstance(candle, dict) and "v" in candle and len(candle["v"]) == 6:
                                ts, o, h, l, c, vol = candle["v"]
                                candle_data = {
                                    "timestamp": int(ts),
                                    "datetime_utc": datetime.fromtimestamp(ts, timezone.utc).strftime('%Y-%m-%d %H:%M:%S'),
                                    "open": o, "high": h, "low": l, "close": c, "volume": vol
                                }
                                chunk_candles.append(candle_data)
                                oldest_timestamp_in_chunk = min(oldest_timestamp_in_chunk, int(ts))

            # Check if the data limit for this request was hit
            if msg.get("m") == "series_completed":
                 if (len(msg.get("p", [])) >= 5 and
                        msg["p"][0] == chart_session_id and
                        msg["p"][1] == series_id_for_replay_data and # Check correct series ID
                        msg["p"][2] == "replay" and
                        isinstance(msg["p"][4], dict) and
                        msg["p"][4].get("data_completed") == "limit"):
                     print("<<< Hit data limit for this request.")
                     hit_limit = True
                     break # Exit receive loop for this chunk

        if not chunk_candles:
            print("No candles received for this chunk, stopping.")
            break

        # Add sorted candles to the main list (optional, could process per chunk)
        chunk_candles.sort(key=lambda x: x["timestamp"])
        all_candles.extend(chunk_candles)
        print(f"Oldest candle in chunk {i+1}: {datetime.fromtimestamp(oldest_timestamp_in_chunk, timezone.utc)}")

        # Prepare timestamp for the *next* chunk (go back from the oldest received)
        # Timeframe is 1 minute (60 seconds)
        current_replay_end_timestamp = oldest_timestamp_in_chunk - 60

        if not hit_limit:
            print("Warning: Did not receive 'data_completed: limit' message. May not have received full chunk.")
            # Decide whether to continue or stop if the limit wasn't hit

        # Optional: Check if oldest_timestamp_in_chunk hasn't changed, meaning no more data available
        # if oldest_timestamp_in_chunk == current_replay_end_timestamp + 60:
        #    print("Oldest timestamp didn't advance. Assuming end of available data.")
        #    break


except Exception as e:
    print(f"An error occurred: {e}")
finally:
    print("\n--- Scraping Finished ---")
    print(f"Total candles scraped: {len(all_candles)}")
    if all_candles:
        all_candles.sort(key=lambda x: x["timestamp"]) # Ensure final sort
        print(f"Earliest candle: {all_candles[0]['datetime_utc']}")
        print(f"Latest candle:   {all_candles[-1]['datetime_utc']}")
        # Example: Print first 5 candles
        print("\nFirst 5 candles:")
        for candle in all_candles[:5]:
            print(candle)
    client.close() # Close WebSocket connection