"""
Multimodal-Mind2Web dataset loader.

Supports loading from HuggingFace (osunlp/Multimodal-Mind2Web) or local files.

Key differences from the original Mind2Web dataset:
  - Each HF row is a **single action step**, not a full trajectory.
  - Screenshots are PIL Image objects (JPEG), not paths or base64 strings.
  - Coordinates come from pos_candidates[].bounding_box_rect (absolute pixels);
    we normalize by the screenshot dimensions.
  - Operation type is stored in operation["op"]: CLICK | TYPE | SELECT | HOVER.
  - Trajectories must be reconstructed by grouping rows on annotation_id.

Available HF splits:
  train         – 7,775 steps / 1,009 tasks
  test_task     – 1,339 steps / 177 tasks  (same websites as training)
  test_website  – 1,019 steps / 142 tasks  (unseen websites)
  test_domain   – 4,060 steps / 694 tasks  (entirely unseen domains)

Usage::

    loader = Mind2WebLoader()
    trajectories = loader.load_trajectories_from_hf(split="train")

Or as a standalone script::

    python -m src.data_pipeline.mind2web_loader \
        --output_dir ./data/mind2web_raw \
        --split train \
        --save_images
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from collections import defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Dataclasses
# ---------------------------------------------------------------------------

@dataclass
class Step:
    """A single annotated step within a web trajectory."""

    action_uid: str
    action_type: str              # click | type | select | scroll
    x: float                     # normalised [0, 1]
    y: float                     # normalised [0, 1]
    value: Optional[str] = None  # text for type/select actions
    screenshot_path: Optional[str] = None  # path after saving PIL Image to disk
    element_id: Optional[str] = None
    action_repr: Optional[str] = None     # human-readable action description
    raw_annotation: Optional[Dict[str, Any]] = field(default=None, repr=False)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "action_uid": self.action_uid,
            "action_type": self.action_type,
            "x": self.x,
            "y": self.y,
            "value": self.value,
            "screenshot_path": self.screenshot_path,
            "element_id": self.element_id,
            "action_repr": self.action_repr,
        }


@dataclass
class Trajectory:
    """A complete annotated web-navigation trajectory."""

    annotation_id: str
    website: str
    domain: str
    task: str
    steps: List[Step] = field(default_factory=list)
    raw: Optional[Dict[str, Any]] = field(default=None, repr=False)

    def __len__(self) -> int:
        return len(self.steps)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "annotation_id": self.annotation_id,
            "website": self.website,
            "domain": self.domain,
            "task": self.task,
            "steps": [s.to_dict() for s in self.steps],
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "Trajectory":
        steps = [
            Step(
                action_uid=s.get("action_uid", ""),
                action_type=s.get("action_type", "click"),
                x=float(s.get("x", 0.5)),
                y=float(s.get("y", 0.5)),
                value=s.get("value"),
                screenshot_path=s.get("screenshot_path"),
                element_id=s.get("element_id"),
                action_repr=s.get("action_repr"),
            )
            for s in d.get("steps", [])
        ]
        return cls(
            annotation_id=d.get("annotation_id", d.get("traj_id", "")),
            website=d.get("website", "unknown"),
            domain=d.get("domain", "unknown"),
            task=d.get("task", ""),
            steps=steps,
        )


# ---------------------------------------------------------------------------
# Action-type mapping
# ---------------------------------------------------------------------------

_ACTION_MAP: Dict[str, str] = {
    "CLICK": "click",
    "TYPE": "type",
    "SELECT": "select",
    "SCROLL": "scroll",
    "HOVER": "click",   # treat hover as click for grounding purposes
    "ENTER": "click",   # treat enter as click
}


def _map_action(raw_type: str) -> str:
    return _ACTION_MAP.get(raw_type.upper(), "click")


# ---------------------------------------------------------------------------
# Coordinate helpers
# ---------------------------------------------------------------------------

def _extract_bbox_rect(candidates: List[Any]) -> Optional[Dict[str, float]]:
    """
    Find the target element's bounding_box_rect from pos_candidates.

    Multimodal-Mind2Web candidates look like:
      {"tag": "button", "backend_node_id": "123",
       "bounding_box_rect": {"x": 100, "y": 200, "width": 50, "height": 30},
       "is_original_target": True, ...}

    We prefer the candidate with is_original_target=True; fall back to the first.
    """
    if not candidates:
        return None

    # Prefer original target
    for c in candidates:
        if isinstance(c, dict) and c.get("is_original_target", False):
            rect = c.get("bounding_box_rect")
            if rect and isinstance(rect, dict):
                return rect

    # Fall back to first candidate with a rect
    for c in candidates:
        if isinstance(c, dict):
            rect = c.get("bounding_box_rect")
            if rect and isinstance(rect, dict):
                return rect

    return None


def _normalize_coords(rect: Dict[str, float], img_w: int, img_h: int) -> Tuple[float, float]:
    """
    Convert an absolute-pixel bounding_box_rect to normalised centre (x, y).

    rect keys: x, y, width, height  (top-left corner + dimensions, in pixels)
    """
    x_px = float(rect.get("x", 0))
    y_px = float(rect.get("y", 0))
    w_px = float(rect.get("width", 0))
    h_px = float(rect.get("height", 0))

    cx = (x_px + w_px / 2.0) / max(img_w, 1)
    cy = (y_px + h_px / 2.0) / max(img_h, 1)

    return max(0.0, min(1.0, cx)), max(0.0, min(1.0, cy))


# ---------------------------------------------------------------------------
# Mind2WebLoader
# ---------------------------------------------------------------------------

class Mind2WebLoader:
    """
    Loads osunlp/Multimodal-Mind2Web from HuggingFace or local storage,
    and reconstructs Trajectory objects by grouping individual action rows.
    """

    HF_DATASET = "osunlp/Multimodal-Mind2Web"

    def __init__(
        self,
        max_steps_per_trajectory: int = 10,
        image_save_dir: Optional[str] = None,
    ) -> None:
        self.max_steps_per_trajectory = max_steps_per_trajectory
        # If set, PIL Images from HF will be saved here as <action_uid>.jpg
        self.image_save_dir = Path(image_save_dir) if image_save_dir else None
        if self.image_save_dir:
            self.image_save_dir.mkdir(parents=True, exist_ok=True)

    # ------------------------------------------------------------------
    # Public: load full trajectories
    # ------------------------------------------------------------------

    def load_trajectories_from_hf(
        self, split: str = "train"
    ) -> List[Trajectory]:
        """
        Download the dataset from HuggingFace using streaming mode to avoid
        loading all 7k+ screenshots into memory at once.

        Rows are processed one-by-one, images saved to disk, then grouped
        by annotation_id into Trajectory objects.

        Args:
            split: One of "train", "test_task", "test_website", "test_domain".
        """
        try:
            import datasets as hf_datasets  # type: ignore
        except ImportError as exc:
            raise ImportError("Install `datasets`: pip install datasets") from exc

        # Raise PIL decompression bomb limit (some screenshots are very large)
        try:
            from PIL import Image as PILImage
            PILImage.MAX_IMAGE_PIXELS = None
        except ImportError:
            pass

        logger.info(
            "Streaming %s split=%r from HuggingFace …", self.HF_DATASET, split
        )

        # streaming=True: rows are fetched one-at-a-time — no memory explosion
        ds = hf_datasets.load_dataset(
            self.HF_DATASET, split=split, streaming=True
        )

        # Collect raw dicts while streaming; images are saved to disk on the fly
        raw_rows: List[Dict[str, Any]] = []
        for i, row in enumerate(ds):
            row_dict = dict(row)
            # Save screenshot immediately and replace the PIL object with a path
            path, w, h = self._handle_screenshot(
                row_dict, str(row_dict.get("action_uid", i))
            )
            row_dict["_screenshot_path"] = path
            row_dict["_img_w"] = w
            row_dict["_img_h"] = h
            # Drop the PIL Image object so it can be garbage-collected
            row_dict.pop("screenshot", None)
            raw_rows.append(row_dict)
            if (i + 1) % 500 == 0:
                logger.info("  … streamed %d rows", i + 1)

        logger.info("Streamed %d rows total.", len(raw_rows))
        return self._group_into_trajectories(raw_rows)

    def load_trajectories_from_local(self, data_dir: str) -> List[Trajectory]:
        """
        Load from a local directory of .jsonl / .json files.
        Expects the same per-step format as the HF dataset (serialised to JSON).
        Screenshots should have already been saved to disk; screenshot_path
        fields in the JSON point to the image files.
        """
        data_dir_path = Path(data_dir)
        if not data_dir_path.exists():
            raise FileNotFoundError(f"Data directory not found: {data_dir}")

        records: List[Dict[str, Any]] = []
        for ext in ("*.jsonl", "*.json"):
            for fp in sorted(data_dir_path.glob(ext)):
                logger.info("Reading %s …", fp)
                with open(fp, "r", encoding="utf-8") as f:
                    for line in f:
                        line = line.strip()
                        if line:
                            records.append(json.loads(line))

        if not records:
            raise FileNotFoundError(f"No data found in {data_dir}")

        logger.info("Loaded %d rows from local files.", len(records))
        return self._group_into_trajectories(records)

    # ------------------------------------------------------------------
    # Grouping rows into trajectories
    # ------------------------------------------------------------------

    def _group_into_trajectories(self, rows) -> List[Trajectory]:
        """
        Group rows (HF dataset rows or dicts) by annotation_id.
        Preserves step order based on target_action_index.
        """
        # bucket rows by annotation_id
        buckets: Dict[str, List[Any]] = defaultdict(list)
        for row in rows:
            aid = str(row.get("annotation_id", "unknown"))
            buckets[aid].append(row)

        trajectories: List[Trajectory] = []
        for annotation_id, step_rows in buckets.items():
            # Sort by target_action_index so steps are in execution order
            step_rows.sort(key=lambda r: int(r.get("target_action_index", 0)))

            # All rows in a bucket share the same website / task / domain
            first = step_rows[0]
            website = str(first.get("website", "unknown"))
            domain = str(first.get("domain", "unknown"))
            task = str(first.get("confirmed_task", first.get("task", "")))

            steps: List[Step] = []
            for row in step_rows:
                if len(steps) >= self.max_steps_per_trajectory:
                    break
                step = self._parse_step(row)
                if step is not None:
                    steps.append(step)

            if steps:  # skip empty trajectories
                trajectories.append(
                    Trajectory(
                        annotation_id=annotation_id,
                        website=website,
                        domain=domain,
                        task=task,
                        steps=steps,
                    )
                )

        logger.info(
            "Grouped into %d trajectories (from %d rows).",
            len(trajectories),
            sum(len(v) for v in buckets.values()),
        )
        return trajectories

    # ------------------------------------------------------------------
    # Parsing a single row into a Step
    # ------------------------------------------------------------------

    def _parse_step(self, row: Any) -> Optional[Step]:
        """Parse one Multimodal-Mind2Web row into a Step dataclass."""
        try:
            row_dict = dict(row) if not isinstance(row, dict) else row

            action_uid = str(row_dict.get("action_uid", ""))

            # ---- action type ----
            # In the HF dataset some fields are JSON-encoded strings; parse them.
            operation = row_dict.get("operation", {}) or {}
            if isinstance(operation, str):
                try:
                    operation = json.loads(operation)
                except Exception:
                    operation = {}
            raw_op = operation.get("op", operation.get("original_op", "CLICK"))
            action_type = _map_action(str(raw_op))

            # ---- value (for type / select) ----
            value: Optional[str] = None
            if action_type in ("type", "select"):
                value = operation.get("value")
                if value is not None:
                    value = str(value)

            # ---- screenshot: use pre-saved path from streaming (preferred) ----
            # In streaming mode, _handle_screenshot() was already called and the
            # PIL object dropped.  Fall back to calling it again for local loads.
            if "_screenshot_path" in row_dict:
                screenshot_path = row_dict["_screenshot_path"]
                img_w = int(row_dict.get("_img_w", 1280))
                img_h = int(row_dict.get("_img_h", 720))
            else:
                screenshot_path, img_w, img_h = self._handle_screenshot(
                    row_dict, action_uid
                )

            # ---- coordinates from pos_candidates ----
            pos_candidates = row_dict.get("pos_candidates", []) or []
            if isinstance(pos_candidates, str):
                try:
                    pos_candidates = json.loads(pos_candidates)
                except Exception:
                    pos_candidates = []
            # Each candidate's bounding_box_rect may also be a JSON string
            parsed_candidates = []
            for c in pos_candidates:
                if isinstance(c, str):
                    try:
                        c = json.loads(c)
                    except Exception:
                        continue
                if isinstance(c, dict):
                    rect_raw = c.get("bounding_box_rect")
                    if isinstance(rect_raw, str):
                        try:
                            c = dict(c)
                            c["bounding_box_rect"] = json.loads(rect_raw)
                        except Exception:
                            pass
                    parsed_candidates.append(c)
            pos_candidates = parsed_candidates

            rect = _extract_bbox_rect(pos_candidates)

            if rect is not None:
                x, y = _normalize_coords(rect, img_w, img_h)
            else:
                # fallback: screen centre
                logger.debug(
                    "No bounding_box_rect for action_uid=%s; using (0.5, 0.5).",
                    action_uid,
                )
                x, y = 0.5, 0.5

            # ---- element id ----
            element_id: Optional[str] = None
            for c in pos_candidates:
                if isinstance(c, dict) and c.get("is_original_target", False):
                    element_id = str(c.get("backend_node_id", ""))
                    break

            # ---- human-readable representation ----
            action_repr: Optional[str] = None
            target_reprs = row_dict.get("target_action_reprs")
            if target_reprs:
                if isinstance(target_reprs, list) and target_reprs:
                    action_repr = str(target_reprs[0])
                else:
                    action_repr = str(target_reprs)

            return Step(
                action_uid=action_uid,
                action_type=action_type,
                x=x,
                y=y,
                value=value,
                screenshot_path=screenshot_path,
                element_id=element_id,
                action_repr=action_repr,
                raw_annotation=row_dict,
            )

        except Exception as exc:
            logger.warning("Failed to parse row: %s", exc)
            return None

    # ------------------------------------------------------------------
    # Screenshot handling
    # ------------------------------------------------------------------

    def _handle_screenshot(
        self, row_dict: Dict[str, Any], action_uid: str
    ) -> Tuple[Optional[str], int, int]:
        """
        Extract the PIL Image from a HF row (or path from a local row),
        optionally save it, and return (path, width, height).

        Returns (None, 1280, 720) if the image cannot be obtained.
        """
        DEFAULT_W, DEFAULT_H = 1280, 720

        # Case 1: already a saved path (local loading)
        if "screenshot_path" in row_dict and row_dict["screenshot_path"]:
            path = str(row_dict["screenshot_path"])
            try:
                from PIL import Image as PILImage  # type: ignore
                with PILImage.open(path) as img:
                    w, h = img.size
                return path, w, h
            except Exception:
                return path, DEFAULT_W, DEFAULT_H

        # Case 2: PIL Image object from HF
        screenshot = row_dict.get("screenshot")
        if screenshot is None:
            return None, DEFAULT_W, DEFAULT_H

        try:
            # HF delivers it as a PIL Image (feature type: Image)
            from PIL import Image as PILImage  # type: ignore
            if isinstance(screenshot, PILImage.Image):
                pil_img = screenshot
            elif isinstance(screenshot, dict):
                # HF sometimes wraps as {"bytes": ..., "path": ...}
                if screenshot.get("bytes"):
                    import io
                    pil_img = PILImage.open(io.BytesIO(screenshot["bytes"]))
                elif screenshot.get("path"):
                    pil_img = PILImage.open(screenshot["path"])
                else:
                    return None, DEFAULT_W, DEFAULT_H
            else:
                return None, DEFAULT_W, DEFAULT_H

            w, h = pil_img.size

            if self.image_save_dir and action_uid:
                save_path = self.image_save_dir / f"{action_uid}.jpg"
                if not save_path.exists():
                    pil_img.convert("RGB").save(save_path, format="JPEG", quality=95)
                return str(save_path), w, h

            return None, w, h

        except Exception as exc:
            logger.warning(
                "Could not process screenshot for action_uid=%s: %s", action_uid, exc
            )
            return None, DEFAULT_W, DEFAULT_H


# ---------------------------------------------------------------------------
# CLI entry-point
# ---------------------------------------------------------------------------

def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Download and convert Multimodal-Mind2Web data"
    )
    parser.add_argument("--output_dir", required=True, help="Where to save data")
    parser.add_argument(
        "--split",
        default="train",
        choices=["train", "test_task", "test_website", "test_domain"],
    )
    parser.add_argument(
        "--local_dir", default=None, help="Load from local dir instead of HF"
    )
    parser.add_argument(
        "--save_images",
        action="store_true",
        help="Save PIL screenshots to <output_dir>/images/",
    )
    parser.add_argument("--max_trajectories", type=int, default=None)
    parser.add_argument("--max_steps", type=int, default=10)
    return parser.parse_args()


def main() -> None:
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
    args = _parse_args()

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    image_save_dir = str(out_dir / "images") if args.save_images else None
    loader = Mind2WebLoader(
        max_steps_per_trajectory=args.max_steps,
        image_save_dir=image_save_dir,
    )

    if args.local_dir:
        trajectories = loader.load_trajectories_from_local(args.local_dir)
    else:
        trajectories = loader.load_trajectories_from_hf(split=args.split)

    if args.max_trajectories:
        trajectories = trajectories[: args.max_trajectories]

    out_file = out_dir / f"mind2web_{args.split}.jsonl"
    with open(out_file, "w", encoding="utf-8") as f:
        for traj in trajectories:
            f.write(json.dumps(traj.to_dict(), ensure_ascii=False) + "\n")

    logger.info(
        "Saved %d trajectories (%d steps total) to %s",
        len(trajectories),
        sum(len(t) for t in trajectories),
        out_file,
    )


if __name__ == "__main__":
    main()