import collections
import logging
import math
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict

import numpy as np
import requests
import tqdm

import json_numpy

json_numpy.patch()


class GO1Client:
    def __init__(self, host: str, port: int):
        self.host = host
        self.port = port
        self.observation_window = None

    def set_language(self, instruction):
        self.instruction = instruction
        print(f"successfully set instruction:{instruction}")

    def predict_action(self, payload: Dict[str, Any]) -> np.ndarray:
        response = requests.post(
            f"http://{self.host}:{self.port}/act", json=payload, headers={"Content-Type": "application/json"}
        )

        if response.status_code == 200:
            result = response.json()
            action = np.array(result)
            return action
        else:
            print(f"Request failed, status code: {response.status_code}")
            print(f"Error message: {response.text}")
            return None

    def update_observation_window(self, input_rgb_arr, input_state):
        self.observation_window = {
            "top": input_rgb_arr[0],
            "left": input_rgb_arr[1],
            "right": input_rgb_arr[2],
            "instruction": self.instruction,
            "state": input_state.reshape(1, -1),
            "ctrl_freqs": np.array([30]),
        }

    def get_action(self):
        assert self.observation_window is not None, "update observation_window first!"
        return self.predict_action(self.observation_window)

    def reset_observation_window(self):
        self.instruction = None
        self.observation_window = None
        print("successfully unset obs and language intruction")


# Encode observation for the model
def encode_obs(observation):
    head_img = observation["observation"]["head_camera"]["rgb"]
    right_img = observation["observation"]["right_camera"]["rgb"]
    left_img = observation["observation"]["left_camera"]["rgb"]

    input_rgb_arr = [head_img, right_img, left_img]
    input_state = observation["joint_action"]["vector"]

    return input_rgb_arr, input_state


def get_model(usr_args):
    return GO1Client(usr_args["host"], usr_args["port"])


def get_payload(TASK_ENV):
    instruction = TASK_ENV.get_instruction()
    observation = TASK_ENV.get_obs()
    input_rgb_arr, input_state = encode_obs(observation)
    return {
        "top": input_rgb_arr[0],
        "left": input_rgb_arr[1],
        "right": input_rgb_arr[2],
        "instruction": str(instruction),
        "state": input_state.reshape(1, -1),
        "ctrl_freqs": np.array([15]),
    }


def eval(TASK_ENV, model, observation):
    if model.observation_window is None:
        instruction = TASK_ENV.get_instruction()
        model.set_language(instruction)

    input_rgb_arr, input_state = encode_obs(observation)
    model.update_observation_window(input_rgb_arr, input_state)

    actions = model.get_action()

    for action in actions:
        TASK_ENV.take_action(action)
        observation = TASK_ENV.get_obs()
        input_rgb_arr, input_state = encode_obs(observation)
        model.update_observation_window(input_rgb_arr, input_state)


def reset_model(model):
    model.reset_observation_window()
