import json
import re
import subprocess
import time
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Union

import numpy as np
import numpy.typing as npt
import requests
from gymnasium import Env
from gymnasium.spaces import Box, Text
from PIL.Image import Image
from playwright.sync_api import CDPSession, Page, Playwright, ViewportSize, expect, sync_playwright

from utils.logger_utils import logger
from utils.timing_utils import time_block, timeit

from .actions import Action, execute_action, get_action_space
from .env_config import set_local_url_mappings
from .processors import ObservationHandler, ObservationMetadata
from .utils import AccessibilityTree, DetachedPage, Observation, get_image_from_url, png_bytes_to_numpy

TIMEOUT = 5000 * 100


# @debug: logging csrf token for debugging purposes
def log_csrf_token(page: Page, msg: str, action: Action | None = None) -> None:
    from .actions import action2str

    try:
        token = page.locator('form.comment-form input[name$="[_token]"]').get_attribute("value", timeout=10)
        if token:
            if action:
                action_str = action2str(action, "som")
                logger.debug(f"{msg}: {token} ACTION: {action_str}")
            else:
                logger.debug(f"{msg}: {token}")

        else:
            if action:
                action_str = action2str(action, "som")
                logger.debug(f"{msg}: no csrf token found ACTION: {action_str}")
            else:
                logger.debug(f"{msg}: no csrf token found")
    except Exception as e:
        pass


@dataclass
class PlaywrightScript:
    function: str  # goto, get_by_role
    destination: str  # https://www.google.com/, combobox
    name: str | None = None  # Search, Avatar 2009
    operation: str | None = None  # click, fill, press
    value: str | None = None  # avatar movie, Enter


def parse_action(action: str) -> PlaywrightScript:
    splitted = action.strip().split(" ")
    assert len(splitted) >= 2
    match splitted[:2]:
        case ["goto", url]:
            assert len(splitted) == 2
            return PlaywrightScript("goto", url)
        case ["get_by_role", destination]:
            assert len(splitted) >= 4
            match splitted[2:]:
                case [name, operation]:
                    return PlaywrightScript("get_by_role", destination, name, operation)
                case [name, operation, value]:
                    return PlaywrightScript("get_by_role", destination, name, operation, value)
                case _:
                    raise ValueError("Invalid action")
        case _:
            raise ValueError(f"Invalid action {action}")


class ScriptBrowserEnv(Env[dict[str, Observation], Action]):
    """
    The goal of this environment is to produce a prototype of a browser environment.
    In the end, we want to support a fully configurable browser environment with wide
    range of action spaces and observation spaces, both structured and unstructured.
    But in this prototype, we just support action space specified by Playwright script,
    and observation space is the html content of the page.
    """

    def __init__(
        self,
        max_page_length: int = 8192,
        headless: bool = True,
        slow_mo: int = 0,
        observation_type: str = "html",
        current_viewport_only: bool = False,
        viewport_size: ViewportSize = {"width": 1280, "height": 720},
        save_trace_enabled: bool = False,
        sleep_after_execution: float = 0.0,
        captioning_fn=None,
        show_scroll_bar: bool = False,
    ):
        # TODO: make Space[Action] = ActionSpace
        self.action_space = get_action_space()  # type: ignore[assignment]
        self.headless = headless
        self.slow_mo = slow_mo
        self.current_viewport_only = current_viewport_only
        self.reset_finished = False
        self.viewport_size = viewport_size
        self.save_trace_enabled = save_trace_enabled
        self.sleep_after_execution = sleep_after_execution
        self.show_scroll_bar = show_scroll_bar
        match observation_type:
            case "html" | "accessibility_tree" | "accessibility_tree_with_captioner":
                self.text_observation_type = observation_type
                self.image_observation_type = ""
                self.main_observation_type = "text"
            case "image":
                self.image_observation_type = observation_type
                self.text_observation_type = ""  # type: ignore[assignment]
                self.main_observation_type = "image"
            case "image_som":
                self.image_observation_type = observation_type
                self.text_observation_type = observation_type  # type: ignore[assignment]
                self.main_observation_type = "image"
            case _:
                raise ValueError(f"Unsupported observation type: {observation_type}")

        self.observation_handler = ObservationHandler(
            self.main_observation_type,
            self.text_observation_type,
            self.image_observation_type,
            self.current_viewport_only,
            self.viewport_size,
            captioning_fn,
        )

        self.observation_space = self.observation_handler.get_observation_space()

    def setup(self, config_file: Path | None = None) -> None:
        self.context_manager = sync_playwright()
        self.playwright = self.context_manager.__enter__()
        ignore_default_args = ["--hide-scrollbars"] if self.show_scroll_bar else []
        self.browser = self.playwright.chromium.launch(
            headless=self.headless, slow_mo=self.slow_mo, ignore_default_args=ignore_default_args
        )

        if config_file:
            instance_config = json.load(config_file.open())

        else:
            instance_config = {}

        storage_state = instance_config.get("storage_state", None)
        start_url = instance_config.get("start_url", None)
        geolocation = instance_config.get("geolocation", None)

        set_local_url_mappings(instance_config["docker_instance_id"])

        # Use custom viewport size if specified in the config, otherwise use the default.
        viewport_size = self.viewport_size.copy()
        viewport_size.update(instance_config.get("viewport_size", {}))
        self.observation_handler.viewport_size = viewport_size

        self.context = self.browser.new_context(
            viewport=viewport_size,
            storage_state=storage_state,
            geolocation=geolocation,
            device_scale_factor=1,
        )
        if self.save_trace_enabled:
            self.context.tracing.start(screenshots=True, snapshots=True)
        if start_url:
            start_urls = start_url.split(" |AND| ")
            for url in start_urls:
                page = self.context.new_page()
                client = page.context.new_cdp_session(page)  # talk to chrome devtools
                if self.text_observation_type in [
                    "accessibility_tree",
                    "accessibility_tree_with_captioner",
                ]:
                    client.send("Accessibility.enable")
                page.client = client  # type: ignore
                page.goto(url, timeout=TIMEOUT)
            # set the first page as the current page
            self.page = self.context.pages[0]
            self.page.bring_to_front()
        else:
            self.page = self.context.new_page()

            client = self.page.context.new_cdp_session(self.page)

            if self.text_observation_type in [
                "accessibility_tree",
                "accessibility_tree_with_captioner",
            ]:
                client.send("Accessibility.enable")
            self.page.client = client  # type: ignore

        # Add reload to help with crsf tokens bug
        self.page.reload()

    def get_page_client(self, page: Page) -> CDPSession:
        return page.client  # type: ignore

    def _get_obs(self) -> dict[str, Observation]:
        obs = self.observation_handler.get_observation(self.page, self.get_page_client(self.page))
        return obs

    def _get_obs_metadata(self) -> dict[str, ObservationMetadata]:
        metadata = self.observation_handler.get_observation_metadata()
        return metadata

    def reset(
        self,
        *,
        seed: int | None = None,
        options: dict[str, str] | None = None,
    ) -> tuple[dict[str, Observation], dict[str, Any]]:
        """
        Reset the environment.
        :param options: options for the environment. The current supported options are:
            - "storage_state": the storage state of the browser. It is a file path to a json file.
        """
        from vwa_utils.vwa_utils import wait_for_page_to_stabilize

        super().reset(seed=seed, options=options)
        if self.reset_finished:
            self.context_manager.__exit__()

        if options is not None and "config_file" in options:
            config_file = Path(options["config_file"])
            if config_file.exists():
                self.setup(config_file=config_file)
            else:
                raise ValueError(f"Config file {config_file} does not exist.")
        else:
            self.setup()
        self.reset_finished = True

        wait_for_page_to_stabilize(
            page=self.page,
            logger=logger,
            min_num_trues=4,
            return_early=False,
            return_after=None,
            hard_sleep=self.sleep_after_execution,
            min_wait_time_seconds=self.sleep_after_execution,
        )

        observation = self._get_obs()
        observation_metadata = self._get_obs_metadata()
        info = {
            "page": DetachedPage(self.page.url, ""),
            "fail_error": "",
            "observation_metadata": observation_metadata,
        }

        return (observation, info)

    def save_trace(self, trace_path: str | Path) -> None:
        if self.save_trace_enabled:
            self.context.tracing.stop(path=trace_path)

    def close(self) -> None:
        if self.reset_finished:
            self.context_manager.__exit__()
        # # Save url2captions
        # with open("url2captions_blip2.json", "w") as wf:
        #     json.dump(self.observation_handler.text_processor.url2caption, wf)
        #     print(f"Saving url2captions_blip2.json with {len(self.observation_handler.text_processor.url2caption)} entries")

    @timeit(custom_name="ENV:step")
    def step(self, action: Action) -> tuple[dict[str, Observation], float, bool, bool, dict[str, Any]]:
        from vwa_utils.vwa_utils import wait_for_page_to_stabilize

        if not self.reset_finished:
            raise RuntimeError("Call reset first before calling step.")

        success = False
        fail_error = ""
        try:
            if "invalid" not in action:
                log_csrf_token(self.page, "csrf token before execute_action", action)
                self.page = execute_action(
                    action,
                    self.page,
                    self.context,
                    self.observation_handler.action_processor,
                )
                success = True
        except Exception as e:
            fail_error = str(e)

        if not success:
            action["wait_for"] = 0.0

        log_csrf_token(self.page, "csrf token before wait_for_page_to_stabilize")

        with time_block("ENV:step:sleep_after_execution"):
            sleep_after_execution = self.sleep_after_execution if "wait_for" not in action else action["wait_for"]
            wait_for_page_to_stabilize(
                page=self.page,
                logger=logger,
                return_early=False,
                min_num_trues=4,
                return_after=None,
                hard_sleep=sleep_after_execution,
                min_wait_time_seconds=sleep_after_execution,
            )

        log_csrf_token(self.page, "csrf token after wait_for_page_to_stabilize")

        observation = self._get_obs()
        observation_metadata = self._get_obs_metadata()

        info = {
            "page": DetachedPage(self.page.url, self.page.content()),
            "fail_error": fail_error,
            "observation_metadata": observation_metadata,
        }
        msg = (
            observation,
            float(success),  # reward
            False,  # terminated
            False,  # truncated
            info,
        )
        return msg
