import base64
import json
import os
import re
from io import BytesIO
from typing import Any

import numpy as np
from bs4 import BeautifulSoup
from PIL import Image

from trajectory_utils.trajectory import VWA_Trajectory
from utils.trajectory_view import TrajectoryView


def parse_webpage_states(html_content, stop_at_critique=False, stop_at_loop_idx=0):
    soup = BeautifulSoup(html_content, "html.parser")
    data = {}
    data["intent_images"] = []
    states = []

    # Extract intent from <pre> tag
    pre_div = soup.find("pre")
    if pre_div:
        # Break by newlines
        pred_div_content = pre_div.text.strip().split("\n")
        for item in pred_div_content:
            if "intent:" in item:
                data["intent"] = item.split(":")[1].strip()
            elif "image:" in item:
                img_content = item.split("image:")[1].strip()
                if img_content.lower() != "none":
                    if img_content.startswith("["):
                        img_urls = eval(img_content)
                        data["intent_images"].extend(img_urls)
                    else:
                        data["intent_images"].append(img_content)

    # Split content by "New Page" headers
    page_sections = soup.find_all("h2")
    crit_loop_idx = 0
    for section in page_sections:
        state = {}
        current_element = section.next_sibling
        skip_elements = []
        while current_element and not (current_element.name == "h2"):
            # 1. URL (found in h3 with class 'url')
            if current_element.name == "h3" and "url" in current_element.get("class", []):
                state["url"] = current_element.find("a").text.replace("URL: ", "")

            if current_element in skip_elements:
                current_element = current_element.next_sibling
                continue

            if current_element.name == "div" and "executor_critique_loop" in current_element.get("class", []):
                skip_elements.append(current_element)

            # 2. State observation and nested elements
            elif current_element.name == "div" and "state_obv" in current_element.get("class", []):
                # Extract action tree
                state["action_tree"] = current_element.find("pre").text if current_element.find("pre") else None

                # Extract image if present
                imgs = current_element.find_all("img")
                state_img = None
                for img in imgs:
                    # Get last b64 image available
                    if img.get("src", "").startswith("data:image/png;base64"):
                        state_img = img["src"]
                state["screenshot"] = state_img

                # Extract previous action
                prev_action = current_element.find("div", class_="prev_action")
                if prev_action:
                    state["prev_action"] = prev_action.text.strip()

                # Extract prediction elements
                predict_div = current_element.find("div", class_="predict_action")
                if predict_div:
                    # Raw parsed prediction
                    raw_pred = predict_div.find("div", class_="raw_parsed_prediction")
                    if raw_pred:
                        state["raw_utterance"] = raw_pred.find("pre").text if raw_pred.find("pre") else None

                    # Parsed action
                    parsed = predict_div.find("div", class_="parsed_action")
                    if parsed:
                        state["parsed_action"] = parsed.find("pre").text if parsed.find("pre") else None

                exec_crit_loop_div = current_element.find("div", class_="executor_critique_loop")
                if stop_at_critique and exec_crit_loop_div:
                    executor_div = current_element.find("div", class_="executor_utterance")
                    if executor_div:
                        raw_utt = executor_div.text.strip()
                        parsed_action = extract_action(raw_utt)

                        if crit_loop_idx == stop_at_loop_idx:
                            state["raw_utterance"] = raw_utt
                            state["parsed_action"] = parsed_action
                        else:
                            state["raw_utterance"] = (
                                f"INITIAL DECISION: {raw_utt}\n\nREVISED DECISION: {state['raw_utterance']}"
                            )
                            state["parsed_action"] = (
                                f"INITIAL ACTION: {parsed_action}\n\nREVISED ACTION: {state['parsed_action']}"
                            )

                        if crit_loop_idx >= stop_at_loop_idx:
                            states.append(state)
                            data["states"] = states
                            return data
                    crit_loop_idx += 1

            current_element = current_element.next_sibling

        if state:  # Only append if we found any data
            states.append(state)

    data["states"] = states
    return data


def extract_action(response: str) -> str:
    # find the first occurence of action
    action_splitter = "```"
    pattern = rf"{action_splitter}((.|\n)*?){action_splitter}"
    match = re.search(pattern, response)
    if match:
        return match.group(1).strip()
    else:
        raise ValueError(f'Cannot find the action identifier "{action_splitter}" in "{response}"')


class Page:
    def __init__(self, url: str):
        self.url = url


def rebuild_trajectory_vwa_format(
    trajectory_data: dict[str, Any] | None = None,
    html_path: str = None,
    stop_at_critique: bool = False,
) -> tuple[list, dict, TrajectoryView]:
    if not trajectory_data and not html_path:
        raise ValueError("Either trajectory_data or html_path must be provided")

    if not trajectory_data:
        trajectory_data = process_html_file(html_path, stop_at_critique)

    trajectory_vwa_format = []
    meta_data = {"action_str_history": ["None"]}
    for state in trajectory_data["states"]:
        observation = {"text": state["action_tree"], "image": state["screenshot"]}

        info = {"url": state["url"], "page": Page(state["url"])}

        extracted_action = ""
        try:
            extracted_action = extract_action(state["raw_utterance"])
        except Exception as e:
            extracted_action = state["parsed_action"]
            print(f"Error extracting action: {e}. Using: {extracted_action}")

        state_vwa_format = {"observation": observation, "info": info}
        action_vwa_format = {"raw_prediction": state["raw_utterance"], "extracted_action": extracted_action}
        meta_data["action_str_history"].append(state["parsed_action"])

        trajectory_vwa_format.append(state_vwa_format)
        trajectory_vwa_format.append(action_vwa_format)
    return trajectory_vwa_format, meta_data, TrajectoryView(trajectory_vwa_format)


def process_html_file(file_path, stop_at_verifier_loop=False, stop_at_loop_idx=0):
    with open(file_path, "r", encoding="utf-8") as f:
        html_content = f.read()

    trajectory_data = parse_webpage_states(html_content, stop_at_verifier_loop, stop_at_loop_idx)
    return trajectory_data


def rebuild_trajectory_vwa_format_from_json(
    trajectory: VWA_Trajectory,
    stop_at_verifier_loop=False,
    stop_at_loop_idx=0,
    ann_types=[],
    verifier_round_stop_idx=0,
):
    trajectory_vwa_format = []
    meta_data = {"action_str_history": ["None"]}

    stop_idx = 0
    for state, action in zip(trajectory.states, trajectory.actions):
        text_obs = state["text"]
        if ann_types:
            trajectory.annotate_images(state=state, actions=action, ann_types=ann_types)
        img_obs = state["images"][-1]

        state_vwa_format = {
            "observation": {"text": text_obs, "image": img_obs},
            "info": {"url": state["url"], "page": Page(state["url"])},
        }

        executor_agent_action = action["agents"]["executor"][-1]
        raw_prediction = executor_agent_action["generated_text"]
        extracted_action = executor_agent_action["parsed_action"]
        action_str_env_parsed = executor_agent_action["action_str_env_parsed"]

        if state.get("scores_per_round", None) is not None and stop_at_verifier_loop:
            if stop_idx == stop_at_loop_idx:
                data = state.get("scores_per_round")[verifier_round_stop_idx]
                raw_prediction = data["raw_prediction"]
                extracted_action = data.get("parsed_action", "")
                if not extracted_action:
                    extracted_action = extract_action(raw_prediction)
                action_str_env_parsed = extracted_action
            stop_idx += 1

        action_vwa_format = {
            "raw_prediction": raw_prediction,
            "extracted_action": extracted_action,
        }
        meta_data["action_str_history"].append(action_str_env_parsed)
        trajectory_vwa_format.append(state_vwa_format)
        trajectory_vwa_format.append(action_vwa_format)

        if stop_idx > stop_at_loop_idx:
            break

    return trajectory_vwa_format, meta_data, TrajectoryView(trajectory_vwa_format)


def extract_trajectory_data(file_path, stop_at_verifier_loop=False, stop_at_loop_idx=0, ann_types=[]):
    if file_path.endswith(".html"):
        execution_data = process_html_file(
            file_path, stop_at_verifier_loop=stop_at_verifier_loop, stop_at_loop_idx=stop_at_loop_idx
        )
        trajectory_data = rebuild_trajectory_vwa_format(execution_data)
        intent = {"text": execution_data["intent"], "images": execution_data["intent_images"]}

        trajectory_view = trajectory_data[2]
        meta_data = trajectory_data[1]
        return intent, trajectory_view, meta_data

    elif file_path.endswith(".json"):
        trajectory = VWA_Trajectory.from_json(file_path)
        trajectory_vwa_format, meta_data, trajectory_view = rebuild_trajectory_vwa_format_from_json(
            trajectory=trajectory,
            stop_at_verifier_loop=stop_at_verifier_loop,
            stop_at_loop_idx=stop_at_loop_idx,
            ann_types=ann_types,
        )
        intent = {"text": trajectory.objective.get("text", ""), "images": trajectory.objective.get("images", [])}

        return intent, trajectory_view, meta_data


if __name__ == "__main__":
    extract_trajectory_data(
        "results/debug/319-319/trajectories/319/trajectory-shopping-319.json",
        stop_at_verifier_loop=True,
        stop_at_loop_idx=2,
        ann_types=["bbox", "coordinates"],
    )
