import base64
import glob
import json
import os
import re
import tempfile
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path
from typing import Any, Dict, List, Optional, TypedDict, Union

import numpy as np
import numpy.typing as npt
import requests
from beartype import beartype
from PIL import Image

from browser_env.env_config import Website, get_website_urls_from_file
from constants.constants import URLS_FILE_TEMPLATE


@dataclass
class DetachedPage:
    url: str
    content: str  # html


@beartype
def png_bytes_to_numpy(png: bytes) -> npt.NDArray[np.uint8]:
    """Convert png bytes to numpy array

    Example:

    >>> fig = go.Figure(go.Scatter(x=[1], y=[1]))
    >>> plt.imshow(png_bytes_to_numpy(fig.to_image('png')))
    """
    return np.array(Image.open(BytesIO(png)))


def get_image_from_url(url: str, headers: dict | None = None, timeout: int = 60) -> Image.Image:
    if not headers:
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
        }
    return Image.open(requests.get(url, stream=True, headers=headers, timeout=timeout).raw)


class DOMNode(TypedDict):
    nodeId: str
    nodeType: str
    nodeName: str
    nodeValue: str
    attributes: str
    backendNodeId: str
    parentId: str
    childIds: list[str]
    cursor: int
    union_bound: list[float] | None
    center: list[float] | None


class AccessibilityTreeNode(TypedDict):
    nodeId: str
    ignored: bool
    role: dict[str, Any]
    chromeRole: dict[str, Any]
    name: dict[str, Any]
    properties: list[dict[str, Any]]
    childIds: list[str]
    parentId: str
    backendDOMNodeId: int
    frameId: str
    bound: list[float] | None
    union_bound: list[float] | None
    offsetrect_bound: list[float] | None
    center: list[float] | None


class BrowserConfig(TypedDict):
    win_upper_bound: float
    win_left_bound: float
    win_width: float
    win_height: float
    win_right_bound: float
    win_lower_bound: float
    device_pixel_ratio: float


class BrowserInfo(TypedDict):
    DOMTree: dict[str, Any]
    config: BrowserConfig


AccessibilityTree = list[AccessibilityTreeNode]
DOMTree = list[DOMNode]

Observation = str | npt.NDArray[np.uint8]


class StateInfo(TypedDict):
    observation: dict[str, Observation]
    info: Dict[str, Any]


def get_task_id_from_config_file(config_file: str | Path) -> str:
    with open(config_file, "r") as f:
        config_file_string = f.read()
    task_config = json.loads(config_file_string)
    return task_config["task_id"]


def load_task_config(
    config_file: str | Path, docker_instance_id: int
) -> tuple[dict[str, Any], str, str, list[Image.Image], list[str]]:
    with open(config_file, "r") as f:
        config_file_string = f.read()
    strings_to_replace = {
        "__SHOPPING__": Website.SHOPPING,
        "__CLASSIFIEDS__": Website.CLASSIFIEDS,
        "__REDDIT__": Website.REDDIT,
        "__WIKIPEDIA__": Website.WIKIPEDIA,
        "__SHOPPING_ADMIN__": Website.SHOPPING_ADMIN,
        "__GITLAB__": Website.GITLAB,
        "__MAP__": Website.MAP,
    }
    websites_urls = get_website_urls_from_file(instance_id=docker_instance_id)
    replace_map = {
        string_to_replace: websites_urls[website]
        for string_to_replace, website in strings_to_replace.items()
        if website in websites_urls
    }
    for key, value in replace_map.items():
        config_file_string = re.sub(re.escape(key), value, config_file_string)
    # TODO: generalize docker separator
    config_file_string = re.sub(r"(?<!\d)_state.json", f"-{docker_instance_id}_state.json", config_file_string)

    task_config = json.loads(config_file_string)

    task_config["websites_urls"] = websites_urls
    intent = task_config["intent"]
    task_id = task_config["task_id"]
    image_paths = task_config.get("image", None)
    intent_images: list[Image.Image] = []
    task_config["docker_instance_id"] = docker_instance_id

    # Regularize image paths data type
    image_paths = [] if image_paths is None else image_paths
    if isinstance(image_paths, str):
        image_paths = [image_paths]

    assert isinstance(image_paths, list), "Image paths must be a list"

    for image_path in image_paths:
        if image_path.startswith("http"):
            input_image = get_image_from_url(image_path)
        else:
            input_image = Image.open(image_path)
        intent_images.append(input_image)

    return task_config, intent, task_id, intent_images, image_paths


def get_instance_id_from_url(url: str) -> str:
    return url.split("-")[-1]


def check_websites(process_id: str, site_names: Optional[List[str]] = None) -> List[Dict[str, str]]:
    """
    Check the status of websites listed in a file, filtered by process_id and/or site_types.
    Returns a list of dicts: { 'site': ..., 'status': 1|0, 'url': ... }
    """
    urls_file = URLS_FILE_TEMPLATE.format(PROCESS_ID=process_id)
    if not os.path.isfile(urls_file):
        print(f"URLs file does not exist: {urls_file}")
        return []

    site_names = site_names or []
    sites_to_check = []
    if process_id:
        if site_names:
            # If both process_id and site_types are provided
            sites_to_check = [f"{site_name}-{process_id}" for site_name in site_names]
        else:
            # If only process_id is provided, filter by ID in the loop
            pass
    results = []
    with open(urls_file, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            parts = line.split()
            if len(parts) < 2:
                continue
            site_id, url = parts[0], parts[1]
            # If process_id is provided but no specific sites, check if site ends with the process_id
            if process_id and not site_names:
                if not site_id.endswith(f"-{process_id}"):
                    continue
            # If specific sites were provided, check against the sites_to_check list
            elif sites_to_check and site_id not in sites_to_check:
                continue
            try:
                resp = requests.get(url, timeout=5)
                status = 1 if resp.ok else 0
            except Exception:
                status = 0
            site_name, instance_id = site_id.split("-")
            results.append({"name": site_name, "status": status, "url": url, "instance_id": instance_id})
    return results


def normalize_url(url: str) -> str:
    # url = re.sub(r"https://", "http://", url)
    # Check if contains `http://`. If not, add it.
    if re.match("^https://", url):
        return url
    if not re.match("http://", url):
        url = "http://" + url
    return url


def map_url_to_local(url: str) -> str:
    """Map real urls to their local counterparts
    Example: https://wikipedia.org -> http://localhost:8888/wikipedia_en_all_maxi_2022-05
    """
    from browser_env.env_config import LOCAL_URLS_NORM, URL_MAPPINGS_FLAT

    url_new = re.sub(r"https://", "http://", url, flags=re.IGNORECASE)
    url_new = re.sub("//en.", "//", url_new, flags=re.IGNORECASE)
    url_new = re.sub("www.", "", url_new, flags=re.IGNORECASE)

    url_modified = False

    for website, local_url, real_url in URL_MAPPINGS_FLAT:
        # relevant only for Wikipedia, which has an additional path for the local url
        local_url = LOCAL_URLS_NORM[website]

        # If real_url counterpart in `url`, map it to the `local` equivalent
        if real_url in url_new:
            url_new = url_new.replace(real_url, local_url)
            url_modified = True
            continue

        real_url_part = real_url.replace("http://", "")
        if re.search(real_url_part, url_new, re.IGNORECASE):
            # local_url_part = local_url.replace("http://", "")
            url_new = url_new.replace(real_url_part, local_url)
            url_new = re.sub("www.", "", url_new)
            url_modified = True
            continue

        # If the mapping is for a `.org` URL, also check for its `.com` variant.
        if real_url.endswith(".org"):
            alt_url = real_url[:-4] + ".com"
            if alt_url in url_new:
                # First, fix the URL so that it uses .org instead of .com.
                url_new = url_new.replace(alt_url, real_url)
                # Now perform the usual mapping.
                url_new.replace(real_url, local_url)
                url_modified = True
                continue

    # If no local urls found, return the original url
    return url_new if url_modified else url


def map_url_to_real(url: str) -> str:
    """Map local urls to their real world counterparts"""
    url = re.sub(r"https://", "http://", url, flags=re.IGNORECASE)
    url = re.sub("www.", "", url, flags=re.IGNORECASE)
    url = re.sub("//en.", "//", url, flags=re.IGNORECASE)

    from browser_env.env_config import LOCAL_URLS_NORM, URL_MAPPINGS_FLAT

    for website, local_url, real_url in URL_MAPPINGS_FLAT:
        local_url_norm = LOCAL_URLS_NORM[website]
        if local_url_norm in url:
            url = url.replace(local_url_norm, real_url)
        elif local_url in url:
            url = url.replace(local_url, real_url)
    return url


def get_site_name(url: str) -> str:
    """Get the site name from the url"""
    from browser_env.env_config import SITE_NAMES, URL_MAPPINGS_FLAT

    for website, local_url, real_url in URL_MAPPINGS_FLAT:
        if local_url in url:
            return SITE_NAMES[website][0]
    return ""


def get_normalized_site_name(url: str) -> tuple[str, str]:
    """Get the normalized site name from the url"""
    from browser_env.env_config import PAGE_NAME_NORMALIZATIONS, URL_MAPPINGS_FLAT

    for website, local_url, real_url in URL_MAPPINGS_FLAT:
        if local_url in url:
            if website in PAGE_NAME_NORMALIZATIONS:
                return PAGE_NAME_NORMALIZATIONS[website]
    return "", ""


def map_endpoint_to_target(url: str, target_endpoint: str = "127.0.0.1") -> str:
    """Map the endpoint of the url to a target counterpart.
    This is needed for String match evals if hosting the benchmark on other machines.
    Example: http://143.215.128.18:8888 -> http://127.0.0.1:8888

    Args:
        url (str): The url to map the endpoint of.
        local_endpoint (str, optional): The local endpoint to map the url to. Defaults to "127.0.0.1".

    Returns:
        str: The url with the endpoint mapped to the local counterpart.
    """
    url = re.sub(r"https://", "http://", url)

    from browser_env.env_config import LOCAL_URLS_NORM, URL_MAPPINGS_FLAT

    for _, instance_url, _ in URL_MAPPINGS_FLAT:
        if instance_url in url:
            base_endpoint = re.sub(r"^https?://", "", instance_url)
            base_endpoint = re.sub(r":[^:]*$", "", base_endpoint)
            return url.replace(base_endpoint, target_endpoint)
    return url


def get_running_websites(process_id: str, sites_to_exclude: list[Website] = []) -> list[Website]:
    all_sites_data = check_websites(process_id=process_id)
    all_sites = [Website(site_data["name"]) for site_data in all_sites_data if site_data["status"] == 1]
    all_sites = [website for website in all_sites if website not in sites_to_exclude]
    return all_sites


if __name__ == "__main__":
    # Get docker instance from cmd line
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--p", type=str, default="", required=False)
    args = parser.parse_args()
    process_id = args.p
    print(get_running_websites(process_id=process_id))
