import json
import random
import re
import time
import tiktoken
from openai import RateLimitError

from mm_agents.vlaa_gui.memory.procedural_memory import PROCEDURAL_MEMORY
import pickle
from PIL import Image
import io
import logging

logger = logging.getLogger("desktopenv.agent")


def get_model_name(agent) -> str:
    engine = getattr(agent, "engine", None)
    if engine is None:
        return "unknown"
    return getattr(engine, "model", engine.__class__.__name__)


def create_pyautogui_code(agent, code: str, obs: dict) -> str:
    """
    Attempts to evaluate the code into a pyautogui code snippet with grounded actions using the observation screenshot.

    Args:
        agent (ACI): The grounding agent to use for evaluation.
        code (str): The code string to evaluate.
        obs (Dict): The current observation containing the screenshot.

    Returns:
        exec_code (str): The pyautogui code to execute the grounded action.

    Raises:
        Exception: If there is an error in evaluating the code.
    """
    agent.assign_screenshot(obs)  # Necessary for grounding
    exec_code = eval(code)
    return exec_code


def call_llm_safe(
    agent,
    temperature: float = 0.0,
    agent_type: str = "unknown",
    subtask: str = "",
    **kwargs,
) -> str:
    from mm_agents.vlaa_gui.core.engine import (
        BedrockRateLimitError,
        LLMRetryExhaustedError,
    )

    max_retries = 5
    attempt = 0
    response = ""
    kwargs.pop("use_thinking", None)
    while attempt < max_retries:
        try:
            response = agent.get_response(temperature=temperature, **kwargs)
            assert response is not None, "Response from agent should not be None"

            break  # If successful, break out of the loop
        except Exception as e:
            attempt += 1
            last_error = e
            model = getattr(agent.engine, "model", "unknown")
            engine_name = type(agent.engine).__name__
            print(
                f"Attempt {attempt} failed [{engine_name} | model={model} | agent_type={agent_type}]: {e}"
            )
            if attempt == max_retries:
                print(
                    f"Max retries reached for {engine_name} | model={model}. Handling failure."
                )
                raise LLMRetryExhaustedError(
                    f"All {max_retries} LLM attempts failed. Last error: {last_error}"
                ) from last_error

            if isinstance(e, (BedrockRateLimitError, RateLimitError)):
                delay = 4.0 * (2 ** (attempt - 1)) + random.uniform(0, 2)
                logger.info(
                    f"Rate limit hit (429). Backing off {delay:.1f}s before retry "
                    f"{attempt + 1}/{max_retries}"
                )
            else:
                delay = 1.0
            time.sleep(delay)
    return response if response is not None else ""


def call_llm_safe_with_thinking(
    agent,
    temperature: float = 0.0,
    agent_type: str = "unknown",
    subtask: str = "",
    **kwargs,
) -> tuple[str, str | None]:
    from mm_agents.vlaa_gui.core.engine import (
        BedrockRateLimitError,
        LLMRetryExhaustedError,
    )

    max_retries = 10
    attempt = 0
    response = ""
    thinking = None
    model_name = get_model_name(agent)
    kwargs.pop("use_thinking", None)
    while attempt < max_retries:
        attempt_number = attempt + 1
        try:
            response, thinking = agent.get_response_with_thinking(
                temperature=temperature, **kwargs
            )
            break
        except Exception as e:
            attempt += 1
            last_error = e
            logger.error(
                "LLM_CALL_EXCEPTION model=%s attempt=%d/%d error=%s",
                model_name,
                attempt_number,
                max_retries,
                str(e),
            )
            print(f"Attempt {attempt} failed: {e}")
            if attempt == max_retries:
                raise LLMRetryExhaustedError(
                    f"All {max_retries} LLM attempts failed. Last error: {last_error}"
                ) from last_error
            if isinstance(e, (BedrockRateLimitError, RateLimitError)):
                delay = 4.0 * (2 ** (attempt - 1)) + random.uniform(0, 2)
                logger.info(
                    f"Rate limit hit (429). Backing off {delay:.1f}s before retry "
                    f"{attempt + 1}/{max_retries}"
                )
            else:
                delay = 1.0
        time.sleep(delay)
    return (response if response is not None else ""), thinking


def call_llm_formatted(
    generator, format_checkers, agent_type: str = "unknown", **kwargs
):
    """
    Calls the generator agent's LLM and ensures correct formatting.

    Args:
        generator (ACI): The generator agent to call.
        obs (Dict): The current observation containing the screenshot.
        format_checkers (Callable): Functions that take the response and return a tuple of (success, feedback).
        agent_type (str): Type of agent for token tracking.
        **kwargs: Additional keyword arguments for the LLM call.

    Returns:
        response (str): The formatted response from the generator agent.
    """
    max_retries = 3  # Set the maximum number of retries
    attempt = 0
    response = ""
    if kwargs.get("messages") is None:
        messages = (
            generator.messages.copy()
        )  # Copy messages to avoid modifying the original
    else:
        messages = kwargs["messages"]
        del kwargs["messages"]  # Remove messages from kwargs to avoid passing it twice
    while attempt < max_retries:
        response = call_llm_safe(
            generator, messages=messages, agent_type=agent_type, **kwargs
        )

        # Prepare feedback messages for incorrect formatting
        feedback_msgs = []
        for format_checker in format_checkers:
            success, feedback = format_checker(response)
            if not success:
                feedback_msgs.append(feedback)
        if not feedback_msgs:
            # logger.info(f"Response formatted correctly on attempt {attempt} for {generator.engine.model}")
            break
        logger.error(
            f"Response formatting error on attempt {attempt} for {generator.engine.model}. Response: {response} {', '.join(feedback_msgs)}"
        )
        messages.append(
            {
                "role": "assistant",
                "content": [{"type": "text", "text": response}],
            }
        )
        logger.info(f"Bad response: {response}")
        delimiter = "\n- "
        formatting_feedback = f"- {delimiter.join(feedback_msgs)}"
        messages.append(
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": PROCEDURAL_MEMORY.FORMATTING_FEEDBACK_PROMPT.replace(
                            "FORMATTING_FEEDBACK", formatting_feedback
                        ),
                    }
                ],
            }
        )
        logger.info("Feedback:\n%s", formatting_feedback)

        attempt += 1
        if attempt == max_retries:
            logger.error(
                "Max retries reached when formatting response. Handling failure."
            )
        time.sleep(1.0)
    return response


def parse_code_from_string(input_string):
    """Parses a string to extract each line of code enclosed in triple backticks (```)

    Args:
        input_string (str): The input string containing code snippets.

    Returns:
        str: The last code snippet found in the input string, or an empty string if no code is found.
    """
    input_string = input_string.strip()

    # This regular expression will match both ```code``` and ```python code```
    # and capture the `code` part. It uses a non-greedy match for the content inside.
    pattern = r"```(?:\w+\s+)?(.*?)```"

    # Find all non-overlapping matches in the string
    matches = re.findall(pattern, input_string, re.DOTALL)
    if len(matches) == 0:
        # return []
        return ""
    relevant_code = matches[
        -1
    ]  # We only care about the last match given it is the grounded action
    return relevant_code


def extract_agent_functions(code):
    """Extracts all agent function calls from the given code.

    Args:
        code (str): The code string to search for agent function calls.

    Returns:
        list: A list of all agent function calls found in the code.
    """
    pattern = r"(agent\.\w+\(\s*.*\))"  # Matches
    return re.findall(pattern, code)


def parse_single_code_from_string(input_string):
    input_string = input_string.strip()
    if input_string.strip() in ["WAIT", "DONE", "FAIL"]:
        return input_string.strip()

    # This regular expression will match both ```code``` and ```python code```
    # and capture the `code` part. It uses a non-greedy match for the content inside.
    pattern = r"```(?:\w+\s+)?(.*?)```"
    # Find all non-overlapping matches in the string
    matches = re.findall(pattern, input_string, re.DOTALL)

    # The regex above captures the content inside the triple backticks.
    # The `re.DOTALL` flag allows the dot `.` to match newline characters as well,
    # so the code inside backticks can span multiple lines.

    # matches now contains all the captured code snippets

    codes = []

    for match in matches:
        match = match.strip()
        commands = [
            "WAIT",
            "DONE",
            "FAIL",
        ]  # fixme: updates this part when we have more commands

        if match in commands:
            codes.append(match.strip())
        elif match.split("\n")[-1] in commands:
            if len(match.split("\n")) > 1:
                codes.append("\n".join(match.split("\n")[:-1]))
            codes.append(match.split("\n")[-1])
        else:
            codes.append(match)

    if len(codes) <= 0:
        return "fail"
    return codes[0]


def resize_screenshot(image_bytes: bytes, new_size: tuple[int, int]) -> bytes:
    image = Image.open(io.BytesIO(image_bytes))
    image_copy = image.copy()
    resized_image = image_copy.resize(new_size)
    output_buffer = io.BytesIO()
    resized_image.save(output_buffer, format=image.format)
    return output_buffer.getvalue()


def get_input_token_length(input_string):
    enc = tiktoken.encoding_for_model("gpt-4")
    tokens = enc.encode(input_string)
    return len(tokens)


def sanitize_code(code):
    # This pattern captures the outermost double-quoted text
    if "\n" in code:
        pattern = r'(".*?")'
        # Find all matches in the text
        matches = re.findall(pattern, code, flags=re.DOTALL)
        if matches:
            # Replace the first occurrence only
            first_match = matches[0]
            code = code.replace(first_match, f'"""{first_match[1:-1]}"""', 1)
    return code


def extract_first_agent_function(code_string):
    # Regular expression pattern to match 'agent' functions with any arguments, including nested parentheses
    pattern = r'agent\.[a-zA-Z_]+\((?:[^()\'"]|\'[^\']*\'|"[^"]*")*\)'

    # Find all matches in the string
    matches = re.findall(pattern, code_string)

    # Return the first match if found, otherwise return None
    return matches[0] if matches else None


def load_knowledge_base(kb_path: str) -> dict:
    try:
        with open(kb_path, "r") as f:
            return json.load(f)
    except Exception as e:
        print(f"Error loading knowledge base: {e}")
        return {}


def load_embeddings(embeddings_path: str) -> dict:
    try:
        with open(embeddings_path, "rb") as f:
            return pickle.load(f)
    except Exception as e:
        print(f"Error loading embeddings: {e}")
        return {}


def save_embeddings(embeddings_path: str, embeddings: dict):
    try:
        with open(embeddings_path, "wb") as f:
            pickle.dump(embeddings, f)
    except Exception as e:
        print(f"Error saving embeddings: {e}")


def split_thinking_response(full_response: str) -> tuple[str, str]:
    try:
        # Extract thoughts section
        thoughts_match = re.search(
            r"<thoughts>(.*?)</thoughts>", full_response, re.DOTALL
        )
        thoughts = thoughts_match.group(1).strip()
        # Extract answer section
        answer_match = re.search(r"<answer>(.*?)</answer>", full_response, re.DOTALL)
        answer = answer_match.group(1).strip()
        return answer, thoughts
    except Exception:
        return full_response, ""
