import struct
import math
import sys
import requests
import os

# --- Configuration ---
SERVER_URL = "http://localhost:8533/"
MAGIC = 72848253210177 # BEPISREA
VERSION = 1
LEVEL_ID = 0

# Physics Constants
DT = 0.02
SPEED = 5.0
GRAVITY = -9.81
JUMP_HEIGHT = 2.0

# Actions
ACTION_NONE = 0
ACTION_FORWARD = 2
ACTION_BACKWARD = 4
ACTION_TAKE_FLAG = 0x20 

def generate_perfect_run(filename="perfect_run.bin"):
    print("[*] Generating TAS Replay...")
    
    # Start at spawn (Standard Unity 0,0,0 usually, but script uses 0, 1.08, 0 based on collider)
    pos = [0.0, 1.08, 0.0] 
    velocity_y = 0.0
    rot = [0.0, 0.0, 0.0, 1.0]
    frames = []

    def add_frames(count, action_mask, set_flag_bit=False):
        nonlocal pos, velocity_y
        for _ in range(count):
            # 1. Input Logic
            move_z = 0.0
            if action_mask & ACTION_FORWARD: move_z += 1.0
            if action_mask & ACTION_BACKWARD: move_z -= 1.0
            
            # 2. Physics Logic (Euler Integration)
            # Y Axis
            is_grounded = pos[1] <= 1.0801
            if is_grounded:
                velocity_y = 0.0
                pos[1] = 1.08
            else:
                velocity_y += GRAVITY * DT
            
            # X/Z Axis
            pos[2] += move_z * SPEED * DT
            pos[1] += velocity_y * DT

            # 3. Record Frame
            # Note: We explicitly set the TAKE_FLAG bit if requested, 
            # to ensure the verifier sees we intended to take it.
            final_action = action_mask
            if set_flag_bit:
                final_action |= ACTION_TAKE_FLAG

            frames.append({
                'pos': list(pos),
                'rot': list(rot),
                'action': final_action
            })

    # --- THE RUN STRATEGY ---
    # 1. Wait at start
    add_frames(10, ACTION_NONE)
    
    # 2. Run Forward (Distance Calculation)
    # 3.5 seconds * 5 units/sec = 17.5 units. 
    # If the flag is at Z=15, this covers it.
    add_frames(175, ACTION_FORWARD) 

    # 3. "Interact" at the destination
    # We walk forward AND simulate the trigger bit for a few frames
    add_frames(10, ACTION_FORWARD, set_flag_bit=True)

    # 4. Run Backward
    add_frames(185, ACTION_BACKWARD)

    # 5. Wait at spawn
    add_frames(10, ACTION_NONE)

    print(f"[*] Generated {len(frames)} frames.")

    # Write File
    with open(filename, "wb") as f:
        f.write(struct.pack('<QHHH', MAGIC, VERSION, LEVEL_ID, len(frames)))
        for frame in frames:
            p = frame['pos']
            r = frame['rot']
            a = frame['action']
            f.write(struct.pack('<fffffffB', 
                p[0], p[1], p[2], 
                r[0], r[1], r[2], r[3], 
                a))
    print(f"[+] Saved TAS run to {filename}")
    return filename

def upload_replay(filename):
    if not os.path.exists(filename):
        print(f"[-] File {filename} not found.")
        return

    print(f"[*] Uploading {filename} to {SERVER_URL}...")
    try:
        with open(filename, "rb") as f:
            data = f.read()
            
        # Send raw bytes
        response = requests.post(SERVER_URL, data=data)
        
        print("\n" + "="*30)
        print(" SERVER RESPONSE")
        print("="*30)
        print(f"Status Code: {response.status_code}")
        print("Body:")
        print(response.text)
        print("="*30 + "\n")
        
        if "lake{" in response.text:
            print("[!!!] FLAG FOUND [!!!]")
    except Exception as e:
        print(f"[-] Error uploading: {e}")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage:")
        print("  python3 replay_solver.py generate   -> Creates and uploads a perfect run")
        print("  python3 replay_solver.py upload <f> -> Uploads an existing replay file")
        sys.exit(1)

    command = sys.argv[1]
    
    if command == "generate":
        f = generate_perfect_run()
        upload_replay(f)
    elif command == "upload":
        if len(sys.argv) < 3:
            print("Please specify a filename.")
        else:
            upload_replay(sys.argv[2])
    else:
        print("Unknown command.")