import json
import logging
import pickle
from enum import Enum
from pathlib import Path
from random import Random
from typing import Any, Dict, List, Literal, Optional, Tuple, cast

import lxml.etree
import requests
from inspect_ai.dataset import Dataset, Sample, hf_dataset
from lxml.etree import _Element
from platformdirs import user_cache_dir
from pydantic import BaseModel
from typing_extensions import TypeAlias

from inspect_evals.mind2web.utils import get_tree_repr, prune_tree

SCORES_URL = "https://buckeyemailosu-my.sharepoint.com/:u:/g/personal/deng_595_buckeyemail_osu_edu/EZllMua3lABAhXQnCN7-pr4BIP4YV8xPfbgyP5FXT18wag?e=yXkK8k&download=1"  # noqa: E501

ActionType = Literal["CLICK", "SELECT", "TYPE"]


class Mind2WebSplit(str, Enum):
    """Enum for Mind2Web dataset splits.

    Each split corresponds to a different generalization setting:
    - CROSS_TASK: Tests generalization to new tasks on websites seen during training
    - CROSS_WEBSITE: Tests generalization to new websites within domains seen during training
    - CROSS_DOMAIN: Tests generalization to entirely new domains and websites
    """

    TRAIN = "train"
    CROSS_TASK = "test_task"
    CROSS_WEBSITE = "test_website"
    CROSS_DOMAIN = "test_domain"


class Answer(BaseModel, frozen=True):
    letter: str
    action: Optional[ActionType] = None
    value: Optional[str] = None


class Mind2WebMetadata(BaseModel, frozen=True):
    final_html: str
    previous_actions: str
    answer: Answer
    target_idx: int
    confirmed_task: str
    task_id: str
    annotation_id: str
    action_uid: str


# Type aliases for utils functions
TreeReprResult: TypeAlias = Tuple[str, Dict[str, int]]

logger = logging.getLogger(__name__)

random = Random()


CACHE_DIR = Path(user_cache_dir("inspect_evals")) / "mind2web" / "data"
EVAL_DIR = Path(__file__).parent


def maybe_download_scores() -> Path:
    """Download scores from SharePoint.

    This function handles the complexity of downloading from SharePoint, which doesn't
    directly serve files via the provided URL. Instead, SharePoint returns an HTML page
    containing the actual download link. We need to parse this HTML to extract the real
    download URL before we can get the actual file.

    Returns:
        Path: The path to the downloaded scores file
    """
    scores_path = CACHE_DIR / "scores.pkl"
    if not scores_path.exists():
        CACHE_DIR.mkdir(parents=True, exist_ok=True)
        response = requests.get(SCORES_URL, allow_redirects=True)
        if response.status_code == 200:
            content_type = response.headers.get("content-type", "")

            logger.info(f"Downloading scores from {SCORES_URL}.")
            logger.info(f"Response content type: {content_type}")
            logger.info(f"Response headers: {dict(response.headers)}")
            # SharePoint doesn't directly serve the file - instead it returns an HTML page
            # with the actual download link embedded in it. We need to parse this HTML to
            # find the real download URL.
            if b"<!DOCTYPE html>" in response.content[:20]:
                logger.info("Received HTML instead of pickle file")
                # Try to find the actual download URL embedded in the SharePoint HTML page
                from bs4 import BeautifulSoup

                soup = BeautifulSoup(response.content, "html.parser")
                download_link = soup.find("a", {"data-download-url": True})
                if download_link:
                    # Extract the real download URL from the SharePoint page
                    download_url = download_link["data-download-url"]  # type: ignore
                    logger.info(f"Found download URL: {download_url}")
                    response = requests.get(download_url, allow_redirects=True)
                    if response.status_code == 200:
                        with open(scores_path, "wb") as f:
                            f.write(response.content)
                    else:
                        raise Exception("Failed to download from extracted URL")
                else:
                    raise Exception("Could not find download URL in SharePoint page")
            else:
                with open(scores_path, "wb") as f:
                    f.write(response.content)
        else:
            raise Exception(
                f"Failed to download problem set from github at {SCORES_URL}."
            )

    assert scores_path.exists()
    return scores_path


def get_dataset(
    limit: int | None = None, split: Mind2WebSplit = Mind2WebSplit.CROSS_DOMAIN
) -> Dataset:
    """Create an Inspect Dataset to run the Mind2Web evaluation.

    This function loads the Mind2Web dataset from the Hugging Face repository
    and prepares it for evaluation. It downloads the necessary score data if not
    already present and formats each sample as a multiple-choice question.

    Args:
        limit: limit number of samples loaded
        split: dataset split

    Returns:
        A Dataset object containing formatted Mind2Web samples.

    Raises:
        Exception: If there is an error loading the score data.
    """
    scores_path = maybe_download_scores()
    # Load and process scores
    score_data = load_scores(scores_path)
    scores = score_data["scores"]
    ranks = score_data["ranks"]

    def record_to_sample(record: dict[str, Any]) -> Sample:
        return format_sample(record, scores, ranks)

    # Load dataset directly
    dataset = hf_dataset(
        "osunlp/Multimodal-Mind2Web",
        limit=limit,
        split=split.value,
        sample_fields=record_to_sample,
    )

    return dataset


def load_scores(score_file_path: Path) -> Dict[str, Dict[str, Any]]:
    """Load and process scores from the pkl file."""
    try:
        logger.info(f"\nAttempting to load scores from: {score_file_path}")
        with open(score_file_path, "rb") as f:
            data = pickle.load(f)
        logger.info("Successfully loaded scores file")
        logger.info(f"Score keys available: {list(data.keys())}")

        # Validate and convert scores/ranks
        scores = data["scores"]
        ranks = data["ranks"]

        # Print sample data for debugging
        if scores:
            sample_id = next(iter(scores))
            logger.debug("\nExample raw data:")
            logger.debug(f"Sample ID: {sample_id}")
            if scores[sample_id]:
                candidate_id = next(iter(scores[sample_id]))
                logger.debug(f"First candidate {candidate_id}:")
                logger.debug(f"Score: {scores[sample_id][candidate_id]}")
                logger.debug(f"Rank: {ranks[sample_id][candidate_id]}")
        logger.info("\nScore data loaded successfully")
        logger.info(f"Number of samples with scores: {len(scores)}")
        if scores:
            first_sample = next(iter(scores))
            logger.debug(f"First sample has {len(scores[first_sample])} candidates")
        return {"scores": scores, "ranks": ranks}
    except FileNotFoundError:
        raise FileNotFoundError(
            f"Score file not found at {score_file_path}. There may have been a problem downloading and caching the scores"  # noqa: E501
        )
    except Exception as e:
        raise Exception(f"Error loading scores file: {e}")


def get_previous_actions(action_reprs: List[str], target_index: int) -> List[str]:
    """Get previous actions up to the target index"""
    return action_reprs[:target_index]


def parse_candidate_json(json_str: str) -> Optional[Dict[str, Any]]:
    """Parse a candidate JSON string and extract node ID."""
    # Fix JSON string formatting
    json_str = json_str.replace('" "', '", "')
    data = json.loads(json_str)

    # Try to get backend_node_id from attributes if it exists
    if "attributes" in data:
        attrs_str = data["attributes"]
        attrs = json.loads(attrs_str)
        node_id = attrs.get("backend_node_id")
    else:
        node_id = data.get("backend_node_id")

    return {"node_id": node_id, "data": data}


def format_sample(
    record: Dict[str, Any],
    scores: Dict[str, Dict[str, float]],
    ranks: Dict[str, Dict[str, int]],
    top_k: int = 5,
) -> Sample:
    """Format a Mind2Web example into a multiple-choice sample."""
    # Get previous actions
    target_index = int(record.get("target_action_index", "0"))
    prev_actions = get_previous_actions(record["action_reprs"], target_index)
    prev_actions_text = "\n".join(prev_actions[-5:]) if prev_actions else "None"

    # Process candidates
    annotation_id = record["annotation_id"]
    action_uid = record["action_uid"]
    sample_id = f"{annotation_id}_{action_uid}"
    candidates = []

    # Process positive candidates
    for pos in record["pos_candidates"]:
        result = parse_candidate_json(pos)
        candidates.append(
            {
                "node_id": result["node_id"],  # type: ignore
                "rank": 0,  # Positive candidates get rank 0
                "is_positive": True,
            }
        )

    # Process negative candidates
    for i, neg in enumerate(record.get("neg_candidates", [])):
        neg_json = parse_candidate_json(neg)
        rank = int(ranks[sample_id][neg_json["node_id"]])  # type: ignore
        candidates.append(
            {
                "node_id": neg_json["node_id"],  # type: ignore
                "rank": rank,
                "is_positive": False,
            }
        )

    # Sort candidates by rank and take top k
    candidates.sort(key=lambda x: x["rank"])
    top_candidates = candidates[:top_k]
    random.shuffle(top_candidates)

    # Process DOM tree
    dom_tree: _Element = lxml.etree.fromstring(record["cleaned_html"])
    dom_tree = cast(  # type: ignore
        _Element,  # type: ignore
        prune_tree(dom_tree, [c["node_id"] for c in top_candidates]),  # type: ignore
    )  # type: ignore
    tree_repr, id_mapping = cast(
        TreeReprResult, get_tree_repr(dom_tree, id_mapping={}, keep_html_brackets=True)
    )

    # Build choices
    choices = ["A. None of the above"]  # Always first choice
    target_index = 0
    target_letter = "A"

    for i, candidate in enumerate(top_candidates):
        letter = chr(66 + i)
        nodes = cast(
            List[_Element],
            dom_tree.xpath(f"//*[@backend_node_id='{candidate['node_id']}']"),
        )
        if nodes:
            node_repr = cast(
                TreeReprResult,
                get_tree_repr(nodes[0], id_mapping=id_mapping, keep_html_brackets=True),
            )[0]
            if candidate["is_positive"]:
                target_index = i + 1
                target_letter = letter
            choices.append(f"{letter}. {node_repr}")

    # Format target and create Answer model if needed
    target = ""
    answer = None
    has_pos_candidate = bool(record["pos_candidates"])
    op = None
    value = None
    if not has_pos_candidate:
        target_index = 0
        target += "A."
    else:
        operation = (
            json.loads(record["operation"])
            if isinstance(record["operation"], str)
            else record["operation"]
        )
        op = operation.get("op", None)
        value = operation.get("value", None)
        target += f"{target_letter}.\nAction: {operation['op']}"
        if value:
            value = value if len(value.strip()) > 0 else None
            target += f"\nValue: {operation['value']}"

    answer = Answer(
        letter=target_letter,
        action=op if op else None,
        value=value,
    )

    # Create metadata model
    metadata = Mind2WebMetadata(
        final_html=tree_repr,
        previous_actions=prev_actions_text,
        answer=answer,
        target_idx=target_index,
        confirmed_task=record["confirmed_task"],
        task_id=annotation_id,
        annotation_id=annotation_id,
        action_uid=action_uid,
    )

    logger.debug(f"Successfully formatted sample\n{metadata}")
    sample = Sample(
        id=sample_id,
        input=sample_id,
        choices=choices,
        target=target,
        metadata=metadata.model_dump(),
    )
    return sample
