import gradio as gr
import os
import numpy as np
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import requests
import json
import uuid
from requests.exceptions import RequestException, HTTPError, Timeout, ConnectionError
from datasets import load_dataset

# Assuming necessary imports are done above

# Global initialization
print("Initializing GPT-2 Tokenizer and Model...")
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')

# Global dataset loading
print("Loading dataset...")
dataset = load_dataset("visionlab/block-towers-10k-3s-trajectory-scale1", split='train')
print("Dataset loaded successfully.")

# Preprocess dataset to build a lookup table for fast access
lookup_table = {}
grid_size = 1.0  # Adjust based on your dataset's scale and distribution
print("Building lookup table...")
for item in dataset:
    # Adjusting to access 'final_positions' as provided in the dataset structure
    for position in item['data']['final_positions']:
        xyz = [position['x'], position['y'], position['z']]
        grid_key = (round(xyz[0]/grid_size), round(xyz[1]/grid_size), round(xyz[2]/grid_size))
        if grid_key not in lookup_table:
            lookup_table[grid_key] = xyz
        else:
            # For simplicity, this keeps the first encountered position for each grid key.
            # You might want to handle multiple positions within the same grid cell differently.
            continue
print("Lookup table built.")

class SecondLifeNavigator:
    def __init__(self):
        self.lookup_table = lookup_table
        print("SecondLifeNavigator initialized with the lookup table.")

    def determine_action_sequence(self, current_xyz):
        print(f"Determining action sequence for position: {current_xyz}")
        grid_key = (round(current_xyz[0]/grid_size), round(current_xyz[1]/grid_size), round(current_xyz[2]/grid_size))
        nearest_position = self.lookup_table.get(grid_key)
        
        if nearest_position:
            action_sequence = f"Move to position {nearest_position[0]}, {nearest_position[1]}, {nearest_position[2]}"
            print(f"Action sequence determined: {action_sequence}")
        else:
            action_sequence = "No suitable position found."
            print(action_sequence)
        return action_sequence

    def send_command_to_corrade(self, corrade_endpoint, command, parameters):
        print(f"Sending command to Corrade: {command}")
        command_data = {
            "command": command,
            "group": "e269893f-a570-0087-930e-6ba2a0b77f9c",
            "password": "nucleus",
        }
        command_data.update(parameters)

        try:
            response = requests.post(corrade_endpoint, json=command_data, timeout=10)
            response.raise_for_status()
            print(f"Command {command} executed successfully.")
            return response.json()
        except Timeout:
            print("Error: The request timed out.")
        except HTTPError as http_err:
            print(f"HTTP error occurred: {http_err}, Status Code: {response.status_code}")
        except ConnectionError:
            print("Error: A connection error occurred.")
        except RequestException as req_err:
            print(f"Error sending command to Corrade: {req_err}")
        return None
            
def create_interface():
    def navigate(current_state):
        endpoint = os.getenv('CORRADE_ENDPOINT', '')  # Ensure this environment variable is correctly set.
        sl_navigator = SecondLifeNavigator()
        # Use determine_action_sequence instead of run
        action_sequence = sl_navigator.determine_action_sequence(current_state)
        random_id = str(uuid.uuid4())
        return random_id + ": " + action_sequence

    interface = gr.Interface(
        fn=navigate,
        inputs=gr.Textbox(lines=2, placeholder="Enter current state..."),
        outputs=gr.Textbox(label="Generated Action Sequence"),
        title="Second Life Navigator",
        description="Generates an action sequence based on the current state for navigating in Second Life.",
    )

    interface.launch()

if __name__ == "__main__":
    create_interface()