import os
import time
import subprocess
from pathlib import Path

def get_correct_coordinates(env, x, y):
    if x is None or y is None:
        # Extract the som informaiton.
        som_info = env.get_som_image(env.render())[0][1]
        x, y = som_info[str(x)][0], som_info[str(x)][1]

        # return "Error: x and y coordinates required for mouse_move"
    else:
        # x = int(x * 1.5)
        # y = int(y * 1.35)
        # x_ = x * 1920 / 1932
        # y_ = y * 1080 / 1092
        x = int(x)
        y = int(y)
    return x, y

# {
#     "type": "function",
#     "function": {
#         "name": "computer_use",
#         "description": "Use a mouse and keyboard to interact with a computer, and take screenshots.\n* This is an interface to a desktop GUI. You do not have access to a terminal or applications menu. You must click on desktop icons to start applications.\n* Some applications may take time to start or process actions, so you may need to wait and take successive screenshots to see the results of your actions. E.g. if you click on Firefox and a window doesn't open, try wait and taking another screenshot.\n* The screen's resolution is 1932x1092.\n* Whenever you intend to move the cursor to click on an element like an icon, you should consult a screenshot to determine the coordinates of the element before moving the cursor.\n* If you tried clicking on a program or link but it failed to load, even after waiting, try adjusting your cursor position so that the tip of the cursor visually falls on the element that you want to click.\n* Make sure to click any buttons, links, icons, etc with the cursor tip in the center of the element. Don't click boxes on their edges unless asked.",
#         "parameters": {
#             "properties": {
#                 "action": {
#                     "description": "The action to perform. The available actions are:\n* `key`: Performs key down presses on the arguments passed in order, then performs key releases in reverse order.\n* `type`: Input a string of text. Use the `clear` parameter to decide whether to overwrite the existing text, and use the `enter` parameter to decide whether the enter key should be pressed after typing the text.\n* `mouse_move`: Move the cursor to a specified (x, y) pixel coordinate on the screen.\n* `click`: Click the left mouse button at a specified (x, y) pixel coordinate on the screen.\n* `drag`: Click at a specified (x, y) pixel coordinate on the screen, and drag the cursor to another specified (x2, y2) pixel coordinate on the screen.\n* `right_click`: Click the right mouse button at a specified (x, y) pixel coordinate on the screen.\n* `middle_click`: Click the middle mouse button at a specified (x, y) pixel coordinate on the screen.\n* `double_click`: Double-click the left mouse button at a specified (x, y) pixel coordinate on the screen.\n* `scroll`: Performs a scroll of the mouse scroll wheel.\n* `wait`: Wait specified seconds for the change to happen.\n* `terminate`: Terminate the current task and report its completion status.",
#                     "enum": [
#                         "key",
#                         "type",
#                         "mouse_move",
#                         "click",
#                         "drag",
#                         "right_click",
#                         "middle_click",
#                         "double_click",
#                         "scroll",
#                         "wait",
#                         "terminate"
#                     ],
#                     "type": "string"
#                 },
#                 "keys": {
#                     "description": "Required only by `action=key`.",
#                     "type": "array"
#                 },
#                 "text": {
#                     "description": "Required only by `action=type`.",
#                     "type": "string"
#                 },
#                 "clear": {
#                     "description": "Assign it to 1 if the text should overwrite the existing text, otherwise assign it to 0. Using this argument clears all text in an element. Required only by `action=type`.",
#                     "type": "number"
#                 },
#                 "enter": {
#                     "description": "Assign it to 1 if the enter key should be pressed after typing the text, otherwise assign it to 0. Required only by `action=type`.",
#                     "type": "number"
#                 },
#                 "coordinate": {
#                     "description": "(x, y): The x (pixels from the left edge) and y (pixels from the top edge) coordinates to move the mouse to.",
#                     "type": "array"
#                 },
#                 "coordinate2": {
#                     "description": "(x2, y2): The x2 (pixels from the left edge) and y2 (pixels from the top edge) coordinates to drag the cursor to. Required only by `action=drag`.",
#                     "type": "array"
#                 },
#                 "pixels": {
#                     "description": "The amount of scrolling to perform. Positive values scroll up, negative values scroll down. This value should be between -5 and 5. Required only by `action=scroll`.",
#                     "type": "number"
#                 },
#                 "time": {
#                     "description": "The seconds to wait. Required only by `action=wait`.",
#                     "type": "number"
#                 },
#                 "status": {
#                     "description": "The status of the task. Required only by `action=terminate`.",
#                     "type": "string",
#                     "enum": [
#                         "success",
#                         "failure"
#                     ]
#                 }
#             },
#             "required": [
#                 "action"
#             ],
#             "type": "object"
#         }
#     }
# }

def computer_control(env, action, coordinate=None, coordinate2=None, text=None, keys=None, clear=False, enter=False, pixels=None, time=None):
    """Handle all mouse/keyboard actions using xdotool"""
    command = action
    if coordinate is not None:
        x, y = coordinate
    if coordinate2 is not None:
        x2, y2 = coordinate2
    try:
        if command == "mouse_move":
            x, y = get_correct_coordinates(env, x, y)
            env.run_command(f"xdotool mousemove --sync {x} {y}")
            return f"Moved mouse to ({x}, {y})"

        elif command == "left_click" or command == "click":
            if x is None and y is None:
                env.run_command("xdotool click 1")
            else:
                x, y = get_correct_coordinates(env, x, y)
                env.run_command(f"xdotool mousemove --sync {x} {y} click 1")

            return "Left clicked"

        elif command == "right_click":
            if x is None and y is None:
                env.run_command("xdotool click 3")
            else:
                x, y = get_correct_coordinates(env, x, y)
                env.run_command(f"xdotool mousemove --sync {x} {y} click 3")
            return "Right clicked"

        elif command == "double_click":
            if x is None and y is None:
                env.run_command("xdotool click --repeat 2 --delay 500 1")
            else:
                x, y = get_correct_coordinates(env, x, y)
                env.run_command(f"xdotool mousemove --sync {x} {y} click --repeat 2 --delay 500 1")
            return "Double clicked"

        elif command == "drag":
            if x is None or y is None:
                return "Error: x and y coordinates required for drag"
            env.run_command(f"xdotool mousemove --sync {x} {y}")
            env.run_command(f"xdotool mousedown 1 mousemove --sync {x2} {y2} mouseup 1")
            return f"Dragged to ({x}, {y})"

        elif command == "type":
            if not text:
                return "Error: text required for typing"
            if clear:
                env.run_command(f'xdotool key ctrl+a')
            env.run_command(f'xdotool type --delay 12 -- "{text}"')
            if enter:
                env.run_command(f'xdotool key Return')
            return f"Typed: {text}"

        elif command == "key":
            if not keys:
                return "Error: keys required for key press"
            new_keys = []
            for key in keys:
                if key == 'enter':
                    key = 'Return'
                new_keys.append(key)
            env.run_command(f"xdotool key {'+'.join(new_keys)}")
            return f"Pressed key: {'+'.join(new_keys)}"

        elif command == "scroll":
            if not pixels:
                return "Error: pixels required for scroll"
            if pixels > 0:
                env.run_command(f"xdotool click 4 --repeat {abs(pixels)}")
            else:
                env.run_command(f"xdotool click 5 --repeat {abs(pixels)}")
            return f"Scrolled: {pixels}"

        elif command == "cursor_position":
            result = env.run_command("xdotool getmouselocation --shell")[0]
            x, y = result.split("\n")[0:2]  # Returns X and Y coordinates
            # x,y = int(x.split("=")[1]), int(y.split("=")[1])
            x = int(x / 1.5)
            y = int(y / 1.35)
            return f"Cursor position: ({x}, {y})"

        elif command == "wait":
            time.sleep(time)
            return f"Waited: {time}"

        return f"Unknown command: {command}"
    except Exception as e:
        return f"Computer control error: {str(e)}"


def screenshot(env, *args, **kwargs):
    """Capture and return screenshot path"""
    try:
        return env.render().resize((1280, 800))
    except Exception as e:
        return f"Screenshot failed: {str(e)}"


def bash(env, command, workdir="/home/devuser/evaluation", root=False):
    """Execute bash command with proper error handling"""
    try:
        if not command:
            return "Error: No command provided"

        result, exit_code = env.run_command(
            command, with_bash_ic=True, workdir=workdir, root=root
        )
        if exit_code != 0:
            return f"Command failed (exit {exit_code}): {result}"
        return result
    except Exception as e:
        return f"Bash error: {str(e)}"


def file_edit(env, command, path, **kwargs):
    """Handle all file operations with validation"""
    try:
        # Validate path format using string operations instead of Path
        if not path.startswith("/"):
            return "Error: Path must be absolute"

        if command == "view":
            return _handle_view(env, path, kwargs.get("view_range"))
        elif command == "create":
            return _handle_create(env, path, kwargs.get("file_text"))
        elif command == "str_replace":
            return _handle_str_replace(
                env, path, kwargs.get("old_str"), kwargs.get("new_str")
            )
        elif command == "insert":
            return _handle_insert(
                env, path, kwargs.get("insert_line"), kwargs.get("new_str")
            )

        return f"Unknown file command: {command}"
    except Exception as e:
        return f"File operation failed: {str(e)}"


def _handle_view(env, path, view_range):
    """View file contents with optional line range"""
    try:
        content = env.file_read(path)
    except Exception as e:
        return f"Error reading file: {str(e)}"

    if content is None:
        return "Error: File does not exist or cannot be read"

    if view_range:
        try:
            start, end = map(int, view_range)
            lines = content.split("\n")[start:end]
            return "\n".join([f"{i+start} | {line}" for i, line in enumerate(lines)])
        except:
            return "Invalid view_range format - use [start_line, end_line]"
    else:
        try:
            lines = content.split("\n")
            return "\n".join([f"{i} | {line}" for i, line in enumerate(lines)])
        except:
            return "Error: File does not exist or cannot be read"


def _handle_create(env, path, content):
    """Create new file with content"""
    if not content:
        return "Error: file_text required for creation"

    try:
        # Attempt to read first to check existence
        existing = env.file_read(path)
        if existing is not None:
            if len(existing.strip()) > 0:
                return "Error: File already exists"
    except:
        pass  # File doesn't exist yet

    try:
        env.file_write(path, content)
        return f"Created file at {path}"
    except Exception as e:
        return f"Error creating file: {str(e)}"


def _handle_str_replace(env, path, old_str, new_str):
    """Replace string in file with validation"""
    if not old_str:
        return "Error: old_str required for replacement"

    try:
        content = env.file_read(path)
    except Exception as e:
        return f"Error reading file: {str(e)}"

    # if content.count(old_str) != 1:
    #     return "Error: old_str must exist exactly once in file"

    new_content = content.replace(old_str, new_str or "")
    try:
        env.file_write(path, new_content)
        return f"Replaced '{old_str}' in {path}"
    except Exception as e:
        return f"Error writing file: {str(e)}"


def _handle_insert(env, path, line, new_str):
    """Insert text at specific line"""
    if line < 0:
        return "Error: Invalid line number"
    if not new_str:
        return "Error: new_str required for insertion"

    try:
        content = env.file_read(path).split("\n")
    except Exception as e:
        return f"Error reading file: {str(e)}"

    if line > len(content):
        return "Error: Line number exceeds file length"

    content.insert(line, new_str)
    try:
        env.file_write(path, "\n".join(content))
        return f"Inserted text at line {line} in {path}"
    except Exception as e:
        return f"Error writing file: {str(e)}"


FUNCTIONS = {
    "computer_control": computer_control,
    "screenshot": screenshot,
    "bash": bash,
    "file_edit": file_edit,
}
