test.py

#!/usr/bin/env python3
"""Export one question as frontend-friendly mock JSON.

Usage:
  python test.py --url https://leetcode.com/problems/two-sum/
  python test.py --url two-sum --out two-sum.mock.json
  python test.py --url two-sum --config 1
"""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path
from typing import Any

import scraper.config as cfg
from scraper.api import create_headers, fetch_playground_codes, fetch_question
from scraper.config import load_config
from scraper.html.assets import load_image_in_b64
from scraper.html.slides import find_slides_json

SLIDE_PRIMARY_RE = re.compile(r"\!\?\!.*?/Documents/.*?\!\?\!", re.IGNORECASE | re.DOTALL)
SLIDE_FALLBACK_1_RE = re.compile(r"\!\?\![^!]*?/Documents/[^!]*?\.json[^!]*?\!\?\!", re.IGNORECASE | re.DOTALL)
SLIDE_FALLBACK_2_RE = re.compile(r"/Documents/[^\s<\"']+\.json", re.IGNORECASE)
PLAYGROUND_IFRAME_RE = re.compile(
    r"<iframe[^>]*src=[\"']([^\"']*playground[^\"']*)[\"'][^>]*>(?:\s*</iframe>)?",
    re.IGNORECASE,
)
MD_IMAGE_RE = re.compile(r"!\[([^\]]*)\]\(([^)]+)\)")
IMG_TAG_RE = re.compile(r"<img\b[^>]*>", re.IGNORECASE)
SRC_ATTR_RE = re.compile(r"\bsrc=[\"']([^\"']+)[\"']", re.IGNORECASE)
ALT_ATTR_RE = re.compile(r"\balt=[\"']([^\"']*)[\"']", re.IGNORECASE)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Export LeetCode question mock JSON for frontend tabs.")
    parser.add_argument("--url", required=True, help="Question URL or titleSlug (e.g. two-sum).")
    parser.add_argument("--out", default=None, help="Output JSON file path (default: ./<slug>.mock.json).")
    parser.add_argument("--config", default=None, help="Config slot number (default: current slot).")
    parser.add_argument(
        "--embed-images-base64",
        action="store_true",
        help="Resolve image URLs and embed as base64 (large output).",
    )
    return parser.parse_args()


def slug_from_url_or_slug(url_or_slug: str) -> str:
    value = url_or_slug.strip()
    if "://" not in value:
        return value.strip("/")
    no_query = value.split("?", 1)[0].split("#", 1)[0].rstrip("/")
    parts = [p for p in no_query.split("/") if p]
    if not parts:
        raise ValueError(f"Cannot parse slug from: {url_or_slug}")
    return parts[-1]


def resolve_image_url(src: str) -> str:
    src = src.strip()
    if not src:
        return src
    if src.startswith("data:"):
        return src
    if src.startswith("//"):
        return f"https:{src}"
    if src.startswith("/"):
        return f"https://leetcode.com{src}"
    split_src = src.split("/")
    if ".." in split_src:
        index = 0
        for i in range(len(split_src) - 1):
            if split_src[i] == ".." and split_src[i + 1] != "..":
                index = i + 1
        return f"https://leetcode.com/explore/{'/'.join(split_src[index:])}"
    return src


def extract_slide_markers(solution_md: str) -> list[str]:
    found = SLIDE_PRIMARY_RE.findall(solution_md)
    if not found:
        found = SLIDE_FALLBACK_1_RE.findall(solution_md)
    if not found:
        found = SLIDE_FALLBACK_2_RE.findall(solution_md)
    return found


def extract_playground_uuid(src_url: str) -> str | None:
    m = re.search(r"/playground/([A-Za-z0-9_-]+)", src_url)
    if m:
        return m.group(1)
    parts = [p for p in src_url.rstrip("/").split("/") if p]
    if len(parts) >= 2:
        return parts[-2]
    return None


def _resolve_image_block(url: str, embed_base64: bool) -> dict[str, Any]:
    resolved = resolve_image_url(url)
    block = {"url": resolved}
    if not embed_base64:
        return block
    if resolved.startswith("data:"):
        block["embedded"] = True
        block["inline_data"] = resolved
        return block
    try:
        block["embedded"] = True
        block["inline_data"] = load_image_in_b64(resolved)
    except Exception as exc:
        block["embedded"] = False
        block["inline_data"] = None
        block["error"] = str(exc)
    return block


def replace_images_with_placeholders(md_text: str, tab_id: str, embed_base64: bool) -> tuple[str, list[dict[str, Any]]]:
    blocks: list[dict[str, Any]] = []
    image_index = 0

    def md_replacer(match: re.Match[str]) -> str:
        nonlocal image_index
        image_index += 1
        alt = match.group(1)
        raw_target = match.group(2).strip()
        unwrapped = raw_target[1:-1] if raw_target.startswith("<") and raw_target.endswith(">") else raw_target
        placeholder = f"[LC_BLOCK:IMAGE:{tab_id}_image_{image_index}]"
        img_block = _resolve_image_block(unwrapped, embed_base64)
        blocks.append(
            {
                "id": f"{tab_id}_image_{image_index}",
                "type": "image",
                "placeholder": placeholder,
                "origin": "markdown_image",
                "alt": alt,
                "original": raw_target,
                **img_block,
            }
        )
        return placeholder

    def html_replacer(match: re.Match[str]) -> str:
        nonlocal image_index
        tag = match.group(0)
        src_match = SRC_ATTR_RE.search(tag)
        if not src_match:
            return tag
        image_index += 1
        raw_src = src_match.group(1)
        alt_match = ALT_ATTR_RE.search(tag)
        alt_text = alt_match.group(1) if alt_match else ""
        placeholder = f"[LC_BLOCK:IMAGE:{tab_id}_image_{image_index}]"
        img_block = _resolve_image_block(raw_src, embed_base64)
        blocks.append(
            {
                "id": f"{tab_id}_image_{image_index}",
                "type": "image",
                "placeholder": placeholder,
                "origin": "html_img_tag",
                "alt": alt_text,
                "original": tag,
                **img_block,
            }
        )
        return placeholder

    out = MD_IMAGE_RE.sub(md_replacer, md_text)
    out = IMG_TAG_RE.sub(html_replacer, out)
    return out, blocks


def sort_blocks_by_position(markdown_text: str, blocks: list[dict[str, Any]]) -> list[dict[str, Any]]:
    for block in blocks:
        block["position"] = markdown_text.find(block["placeholder"])
    blocks.sort(key=lambda x: x["position"])
    return blocks


def main() -> None:
    args = parse_args()
    if args.config is not None:
        cfg.selected_config = str(args.config)

    conf = load_config()
    headers = create_headers(conf.leetcode_cookie)
    slug = slug_from_url_or_slug(args.url)

    question = fetch_question(headers, slug)
    question_url = f"https://leetcode.com/problems/{slug}/"
    solution_md_original = (question.get("solution") or {}).get("content") or "No Solution"
    question_md_original = question.get("content") or ""
    hints = question.get("hints") or []
    hints_md_original = "\n".join(f"{i + 1}. {h}" for i, h in enumerate(hints)) if hints else "No Hints"

    question_md, question_image_blocks = replace_images_with_placeholders(
        question_md_original,
        tab_id="question",
        embed_base64=args.embed_images_base64,
    )
    hints_md, hints_image_blocks = replace_images_with_placeholders(
        hints_md_original,
        tab_id="hints",
        embed_base64=args.embed_images_base64,
    )
    solution_md = solution_md_original

    slide_markers = extract_slide_markers(solution_md_original)
    slide_timelines = find_slides_json(solution_md_original)
    slide_blocks: list[dict[str, Any]] = []
    for i, marker in enumerate(slide_markers, start=1):
        timeline = slide_timelines[i - 1] if i - 1 < len(slide_timelines) else []
        placeholder = f"[LC_BLOCK:SLIDES:slides_{i}]"
        solution_md = solution_md.replace(marker, placeholder, 1)
        resolved_timeline = []
        for frame in timeline:
            f = dict(frame)
            if f.get("image"):
                f["image_original"] = f["image"]
                img_block = _resolve_image_block(f["image"], args.embed_images_base64)
                f["image"] = img_block["url"]
                if "inline_data" in img_block:
                    f["image_inline_data"] = img_block.get("inline_data")
                    f["image_embedded"] = img_block.get("embedded", False)
                if "error" in img_block:
                    f["image_error"] = img_block["error"]
            resolved_timeline.append(f)
        slide_blocks.append(
            {
                "id": f"slides_{i}",
                "type": "slides",
                "placeholder": placeholder,
                "source_marker": marker,
                "timeline": resolved_timeline,
            }
        )

    playground_blocks: list[dict[str, Any]] = []
    for i, match in enumerate(list(PLAYGROUND_IFRAME_RE.finditer(solution_md_original)), start=1):
        full_iframe = match.group(0)
        src_url = match.group(1)
        uuid = extract_playground_uuid(src_url)
        codes = fetch_playground_codes(headers, uuid) if uuid else []
        placeholder = f"[LC_BLOCK:PLAYGROUND:playground_{i}]"
        solution_md = solution_md.replace(full_iframe, placeholder, 1)
        playground_blocks.append(
            {
                "id": f"playground_{i}",
                "type": "playground",
                "placeholder": placeholder,
                "iframe_src": src_url,
                "uuid": uuid,
                "codes": codes,
            }
        )

    solution_md, solution_image_blocks = replace_images_with_placeholders(
        solution_md,
        tab_id="solution",
        embed_base64=args.embed_images_base64,
    )

    question_blocks = sort_blocks_by_position(question_md, question_image_blocks)
    hints_blocks = sort_blocks_by_position(hints_md, hints_image_blocks)
    solution_blocks = sort_blocks_by_position(solution_md, slide_blocks + playground_blocks + solution_image_blocks)

    output = {
        "schemaVersion": 1,
        "source": {
            "questionUrl": question_url,
            "titleSlug": slug,
            "exportedBy": "test.py",
        },
        "question": {
            "questionId": question.get("questionId"),
            "title": question.get("title"),
            "titleSlug": slug,
            "difficulty": question.get("difficulty"),
            "submitUrl": question.get("submitUrl"),
            "exampleTestcaseList": question.get("exampleTestcaseList"),
            "similarQuestions": json.loads(question.get("similarQuestions") or "[]"),
            "companyTagStatsV2": json.loads(question.get("companyTagStatsV2") or "{}"),
            "defaultCodeByLang": json.loads(question.get("codeDefinition") or "[]"),
        },
        "tabs": [
            {
                "id": "question",
                "title": "Question",                
                "contentType": "markdown_with_blocks",
                "markdownRaw": question_md_original,
                "markdownResolved": question_md,
                "blocks": question_blocks,
            },
            {
                "id": "hints",
                "title": "Hints",
                "contentType": "markdown_with_blocks",
                "markdownRaw": hints_md_original,
                "markdownResolved": hints_md,
                "blocks": hints_blocks,
            },
            {
                "id": "solution",
                "title": "Solution",
                "contentType": "markdown_with_blocks",
                "markdownRaw": solution_md_original,
                "markdownResolved": solution_md,
                "blocks": solution_blocks,
            },
        ],
        "apiRaw": {
            "GetQuestion": question,
        },
    }

    out_path = Path(args.out) if args.out else Path(f"{slug}.mock.json")
    out_path.write_text(json.dumps(output, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"Saved mock JSON: {out_path.resolve()}")
    print(
        "Blocks => "
        f"question:{len(question_blocks)} "
        f"hints:{len(hints_blocks)} "
        f"solution:{len(solution_blocks)} "
        f"(slides:{len(slide_blocks)}, playgrounds:{len(playground_blocks)}, "
        f"solution_images:{len(solution_image_blocks)})"
    )


if __name__ == "__main__":
    main()