#!/usr/bin/env python
"""Inference entrypoint for rich Chinese CMGUI screenshot summarization.

The script supports two modes:

1. Model inference from a trained rich checkpoint.
2. OCR-template baseline without a checkpoint.

The template mode is intentionally simple. It exists for debugging and as a
report baseline, not as the final system.
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple

import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

from train_rich import (  # noqa: E402
    RichCollator,
    RichScreenshotDataset,
    apply_structured_evidence_predictions,
    apply_structured_function_predictions,
    extract_summary,
    load_rich_checkpoint,
    move_batch,
    natural_prediction_from_text,
    prediction_from_summary,
    repair_prediction_with_context,
    safe_json_loads,
    target_schema_is_natural_text,
    target_schema_is_summary,
)


DEFAULT_TEST_FILE = "data/rich_cmgui/processed/test_rich_teacher500_natural_qwen8000.jsonl"
DEFAULT_OUTPUT_JSONL = "outputs/rich_infer_current/summary.jsonl"
DEFAULT_OUTPUT_MD = "outputs/rich_infer_current/summary.md"
DEFAULT_FUNCTION_THRESHOLD = 0.20
DEFAULT_SEARCH_THRESHOLD = 0.20
DEFAULT_EVIDENCE_THRESHOLD = 0.50
DEFAULT_MAX_STRUCTURED_ITEMS = 8


def str_to_bool(value: str) -> bool:
    return str(value).lower() in {"1", "true", "yes"}


def find_latest_rich_checkpoint(runs_dir: str | Path = "runs") -> Optional[Path]:
    """Find the newest usable stage3/stage4 rich checkpoint for demo/inference."""
    root = Path(runs_dir)
    if not root.exists():
        return None
    candidates: List[Tuple[float, int, Path]] = []
    for checkpoint in root.rglob("checkpoint-best"):
        config_path = checkpoint / "rich_config.json"
        if not config_path.exists():
            continue
        normalized = checkpoint.as_posix().lower()
        if "stage3_" in normalized:
            stage_rank = 3
        elif "stage4_" in normalized:
            stage_rank = 4
        else:
            continue
        try:
            mtime = max(checkpoint.stat().st_mtime, config_path.stat().st_mtime)
        except OSError:
            continue
        candidates.append((mtime, stage_rank, checkpoint))
    if not candidates:
        return None
    candidates.sort(key=lambda item: (item[0], item[1]), reverse=True)
    return candidates[0][2]


SUMMARY_KEYS = ["画面总结", "summary_zh", "summary"]
VISIBLE_KEYS = ["可见文字", "visible_text"]
INTERACTION_KEYS = ["互动数据", "interaction_data"]
FUNCTION_KEYS = ["功能入口", "ui_functions"]
EVIDENCE_KEYS = ["关键证据", "key_ui_clues", "evidence"]

FUNCTION_KEYWORDS = [
    "搜索",
    "分享",
    "评论",
    "点赞",
    "收藏",
    "关注",
    "返回",
    "首页",
    "消息",
    "拍同款",
    "购物车",
    "购买",
    "下单",
    "展开",
    "关闭",
    "保存",
]

INTERACTION_PATTERNS = [
    ("点赞数", re.compile(r"(?:点赞|赞|like)\s*[:：]?\s*([0-9]+(?:\.[0-9]+)?[万kK]?)")),
    ("评论数", re.compile(r"(?:评论|comment)\s*[:：]?\s*([0-9]+(?:\.[0-9]+)?[万kK]?)")),
    ("收藏数", re.compile(r"(?:收藏|star|favorite)\s*[:：]?\s*([0-9]+(?:\.[0-9]+)?[万kK]?)")),
    ("分享数", re.compile(r"(?:分享|share)\s*[:：]?\s*([0-9]+(?:\.[0-9]+)?[万kK]?)")),
    ("未读消息", re.compile(r"(?:消息|未读)\s*[:：]?\s*([0-9]+(?:\.[0-9]+)?[万kK]?)")),
]

COMMON_NOISE = {
    "Q",
    "0",
    "1",
    "2",
    "3",
    "4",
    "5",
    "6",
    "7",
    "8",
    "9",
}


def read_jsonl(path: Path) -> Iterable[Dict[str, Any]]:
    with path.open("r", encoding="utf-8") as f:
        for line_no, line in enumerate(f, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                yield json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"Bad JSON at {path}:{line_no}: {exc}") from exc


def write_jsonl(path: Path, rows: Iterable[Dict[str, Any]]) -> int:
    path.parent.mkdir(parents=True, exist_ok=True)
    count = 0
    with path.open("w", encoding="utf-8", newline="\n") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")
            count += 1
    return count


def safe_text(value: Any) -> str:
    if value is None:
        return ""
    return re.sub(r"\s+", " ", str(value)).strip()


def first_value(obj: Dict[str, Any], keys: List[str], default: Any = None) -> Any:
    for key in keys:
        if key in obj and obj[key] not in (None, ""):
            return obj[key]
    return default


def normalize_prediction(obj: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
    if not isinstance(obj, dict):
        return None
    return {
        "画面总结": safe_text(first_value(obj, SUMMARY_KEYS, "")),
        "可见文字": first_value(obj, VISIBLE_KEYS, []) or [],
        "互动数据": first_value(obj, INTERACTION_KEYS, []) or [],
        "功能入口": first_value(obj, FUNCTION_KEYS, []) or [],
        "关键证据": first_value(obj, EVIDENCE_KEYS, []) or [],
    }


def item_id(item: Dict[str, Any]) -> str:
    return safe_text(item.get("id") or item.get("ocr_id"))


def build_item_index(row: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
    index: Dict[str, Dict[str, Any]] = {}
    for item in row.get("ui_items", []) or []:
        if item.get("id"):
            index[str(item["id"])] = item
        if item.get("ocr_id"):
            index[str(item["ocr_id"])] = item
    for item in row.get("ocr_items", []) or []:
        if item.get("id") and str(item["id"]) not in index:
            index[str(item["id"])] = item
    return index


def is_noise_text(text: str) -> bool:
    if not text:
        return True
    if text in COMMON_NOISE:
        return True
    if re.fullmatch(r"[0-9:：/.\-\s]+", text) and len(text) <= 6:
        return True
    if len(text) == 1 and not re.search(r"[\u4e00-\u9fffA-Za-z]", text):
        return True
    return False


def collect_visible_text(row: Dict[str, Any], max_items: int) -> List[str]:
    candidates: List[Tuple[float, str]] = []
    for item in row.get("ocr_items", []) or []:
        text = safe_text(item.get("text"))
        if is_noise_text(text):
            continue
        conf = float(item.get("conf", item.get("ocr_conf", 1.0)) or 1.0)
        if conf < 0.55:
            continue
        score = conf + min(len(text), 30) / 100.0
        if any(k in text for k in FUNCTION_KEYWORDS):
            score += 0.15
        candidates.append((score, text[:80]))
    seen = set()
    visible: List[str] = []
    for _, text in sorted(candidates, key=lambda x: x[0], reverse=True):
        if text not in seen:
            visible.append(text)
            seen.add(text)
        if len(visible) >= max_items:
            break
    return visible


def collect_functions(row: Dict[str, Any], max_items: int) -> List[Dict[str, Any]]:
    funcs: List[Dict[str, Any]] = []
    seen = set()
    for item in row.get("ui_items", []) or []:
        text = safe_text(item.get("text"))
        if not text:
            continue
        for keyword in FUNCTION_KEYWORDS:
            if keyword in text and keyword not in seen:
                funcs.append({"name": keyword, "evidence_ids": [item_id(item)] if item_id(item) else []})
                seen.add(keyword)
                break
        if len(funcs) >= max_items:
            break
    return funcs


def collect_interactions(row: Dict[str, Any], max_items: int) -> List[Dict[str, Any]]:
    interactions: List[Dict[str, Any]] = []
    seen = set()
    for item in row.get("ui_items", []) or []:
        text = safe_text(item.get("text"))
        if not text:
            continue
        for name, pattern in INTERACTION_PATTERNS:
            match = pattern.search(text)
            if match and (name, match.group(1)) not in seen:
                interactions.append(
                    {
                        "name": name,
                        "value": match.group(1),
                        "evidence_ids": [item_id(item)] if item_id(item) else [],
                    }
                )
                seen.add((name, match.group(1)))
                break
        if len(interactions) >= max_items:
            break
    return interactions


def template_prediction(row: Dict[str, Any], max_visible: int = 12) -> Dict[str, Any]:
    app = safe_text(row.get("app")) or "移动应用"
    visible = collect_visible_text(row, max_visible)
    funcs = collect_functions(row, max_items=8)
    interactions = collect_interactions(row, max_items=8)
    content_hint = "、".join(visible[:3]) if visible else "当前页面内容"
    func_hint = "、".join(f["name"] for f in funcs[:4]) if funcs else "查看和操作"
    instruction = safe_text(row.get("instruction"))
    if instruction:
        summary = f"这是一个{app}页面，屏幕主要展示{content_hint}等内容，可通过{func_hint}等入口继续操作；当前任务语境是{instruction}。"
    else:
        summary = f"这是一个{app}页面，屏幕主要展示{content_hint}等内容，可通过{func_hint}等入口继续操作。"
    evidence: List[str] = []
    for value in row.get("weak_evidence_ids", []) or []:
        if value not in evidence:
            evidence.append(value)
    for func in funcs:
        for value in func.get("evidence_ids", []) or []:
            if value and value not in evidence:
                evidence.append(value)
    for item in row.get("ui_items", [])[: max(0, 8 - len(evidence))]:
        value = item_id(item)
        if value and value not in evidence:
            evidence.append(value)
    return {
        "画面总结": summary,
        "可见文字": visible,
        "互动数据": interactions,
        "功能入口": funcs,
        "关键证据": evidence[:8],
    }


def clue_details(row: Dict[str, Any], clue_ids: List[str], score_map: Optional[Dict[str, float]] = None) -> List[Dict[str, Any]]:
    index = build_item_index(row)
    out: List[Dict[str, Any]] = []
    for clue_id in clue_ids:
        clue_id = safe_text(clue_id)
        item = index.get(clue_id)
        if not item:
            continue
        out.append(
            {
                "element_id": clue_id,
                "type": item.get("type"),
                "text": item.get("text", ""),
                "bbox": item.get("bbox"),
                "location": item.get("location"),
                "source": item.get("source"),
                "score": round(float(score_map.get(clue_id, 0.0)), 4) if score_map else None,
            }
        )
    return out


def row_result(
    row: Dict[str, Any],
    raw_text: str,
    pred_obj: Optional[Dict[str, Any]],
    json_valid: bool,
    evidence_scores: Optional[Dict[str, float]],
    allow_template_fallback: bool,
    source: str,
) -> Dict[str, Any]:
    structured = normalize_prediction(pred_obj)
    used_fallback = False
    if structured is None or not structured.get("画面总结"):
        if allow_template_fallback:
            structured = template_prediction(row)
            used_fallback = True
        else:
            structured = {
                "画面总结": "",
                "可见文字": [],
                "互动数据": [],
                "功能入口": [],
                "关键证据": [],
            }
    clue_ids = [safe_text(x) for x in structured.get("关键证据", []) if safe_text(x)]
    return {
        "screen_id": row.get("screen_id"),
        "image_path": row.get("image_path"),
        "app": row.get("app"),
        "summary": structured.get("画面总结", ""),
        "prediction": structured,
        "prediction_raw": raw_text,
        "json_valid": bool(json_valid and not used_fallback),
        "used_template_fallback": used_fallback,
        "source": "template" if used_fallback else source,
        "key_ui_clues": clue_details(row, clue_ids, evidence_scores),
    }


def run_template(args: argparse.Namespace) -> List[Dict[str, Any]]:
    rows = list(read_jsonl(Path(args.input_file)))
    if args.max_samples:
        rows = rows[: args.max_samples]
    results = []
    for row in tqdm(rows, desc="template infer"):
        pred = template_prediction(row, max_visible=args.max_visible_text)
        results.append(
            row_result(
                row=row,
                raw_text=json.dumps(pred, ensure_ascii=False),
                pred_obj=pred,
                json_valid=True,
                evidence_scores=None,
                allow_template_fallback=False,
                source="template",
            )
        )
    return results


@torch.no_grad()
def run_model(args: argparse.Namespace) -> List[Dict[str, Any]]:
    device = torch.device(args.device if args.device else ("cuda" if torch.cuda.is_available() else "cpu"))
    model, tokenizer, image_processor, ckpt_args = load_rich_checkpoint(args.checkpoint, device)
    ckpt_args.num_workers = args.num_workers
    if args.generation_no_repeat_ngram_size is not None:
        ckpt_args.generation_no_repeat_ngram_size = args.generation_no_repeat_ngram_size
    if args.generation_repetition_penalty is not None:
        ckpt_args.generation_repetition_penalty = args.generation_repetition_penalty
    if args.generation_block_extra_ids is not None:
        ckpt_args.generation_block_extra_ids = args.generation_block_extra_ids
    if args.generation_block_title_prefix is not None:
        ckpt_args.generation_block_title_prefix = args.generation_block_title_prefix
    if args.generation_force_json_start is not None:
        ckpt_args.generation_force_json_start = args.generation_force_json_start
    runtime_names = [
        "canonicalize_targets",
        "drop_bare_search_functions",
        "structured_function_mode",
        "structured_function_threshold",
        "structured_search_threshold",
        "structured_max_functions",
        "structured_strict_search_candidates",
        "structured_evidence_mode",
        "structured_evidence_threshold",
        "structured_max_evidence",
        "structured_evidence_fallback_top1",
    ]
    for name in runtime_names:
        value = getattr(args, name, None)
        if value is not None and value != "":
            setattr(ckpt_args, name, value)
    dataset = RichScreenshotDataset(args.input_file, max_samples=args.max_samples)
    collator = RichCollator(tokenizer, image_processor, ckpt_args)
    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=collator,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
    )
    results: List[Dict[str, Any]] = []
    model.eval()
    for batch in tqdm(loader, desc="model infer"):
        rows = batch["rows"]
        batch = move_batch(batch, device)
        texts = model.generate_text(batch, tokenizer, num_beams=args.num_beams, max_new_tokens=args.max_new_tokens)
        _, _, elem_tokens, elem_key_padding = model.build_memory(batch)
        evidence_scores_tensor = torch.sigmoid(model.evidence_head(elem_tokens).squeeze(-1)).detach().cpu()
        function_scores_tensor = torch.sigmoid(model.ui_function_head(elem_tokens).squeeze(-1)).detach().cpu()
        search_scores_tensor = torch.sigmoid(model.search_function_head(elem_tokens).squeeze(-1)).detach().cpu()
        scores = evidence_scores_tensor.numpy()
        masks = (~elem_key_padding).detach().cpu().numpy()
        for row_idx, (row, text) in enumerate(zip(rows, texts)):
            if target_schema_is_summary(getattr(ckpt_args, "target_schema", "zh")):
                pred_obj = prediction_from_summary(row, text)
                ok = True
            elif target_schema_is_natural_text(getattr(ckpt_args, "target_schema", "zh")):
                pred_obj = natural_prediction_from_text(text)
                ok = bool(extract_summary(pred_obj))
            else:
                pred_obj, ok = safe_json_loads(text)
            evidence_scores: Dict[str, float] = {}
            elements = row.get("ui_items", []) or []
            for elem_idx, elem in enumerate(elements[: scores.shape[1]]):
                if masks[row_idx, elem_idx]:
                    elem_id = item_id(elem)
                    if elem_id:
                        evidence_scores[elem_id] = float(scores[row_idx, elem_idx])
            if pred_obj is None:
                top_ids = [
                    item_id(elements[i])
                    for i in np.argsort(scores[row_idx])[::-1]
                    if i < len(elements) and masks[row_idx, i] and item_id(elements[i])
                ][: args.top_k_clues]
                pred_obj = {"关键证据": top_ids}
            repair_applied = False
            if args.context_summary_repair:
                pred_obj, repair_applied = repair_prediction_with_context(row, pred_obj)
                ok = True
            pred_obj = apply_structured_function_predictions(
                row,
                pred_obj,
                function_scores_tensor[row_idx],
                search_scores_tensor[row_idx],
                ckpt_args,
            )
            pred_obj = apply_structured_evidence_predictions(row, pred_obj, evidence_scores_tensor[row_idx], ckpt_args)
            results.append(
                row_result(
                    row=row,
                    raw_text=text,
                    pred_obj=pred_obj,
                    json_valid=ok,
                    evidence_scores=evidence_scores,
                    allow_template_fallback=args.allow_template_fallback,
                    source="model+context_repair" if repair_applied else "model",
                )
            )
    return results


def markdown_for_result(row: Dict[str, Any]) -> str:
    pred = row.get("prediction") or {}
    lines = [
        f"## {row.get('screen_id')}",
        "",
        "### 画面总结",
        safe_text(pred.get("画面总结")),
        "",
    ]
    visible = pred.get("可见文字") or []
    if visible:
        lines += ["### 可见文字 / 文案"]
        lines += [f"- {safe_text(x)}" for x in visible[:16]]
        lines.append("")
    interactions = pred.get("互动数据") or []
    if interactions:
        lines += ["### 互动数据"]
        for item in interactions:
            if isinstance(item, dict):
                value = safe_text(item.get("value"))
                suffix = f": {value}" if value else ""
                lines.append(f"- {safe_text(item.get('name'))}{suffix}")
        lines.append("")
    funcs = pred.get("功能入口") or []
    if funcs:
        lines += ["### 功能入口"]
        for item in funcs:
            if isinstance(item, dict):
                evidence = ", ".join(item.get("evidence_ids", []) or [])
                suffix = f" ({evidence})" if evidence else ""
                lines.append(f"- {safe_text(item.get('name'))}{suffix}")
        lines.append("")
    clues = row.get("key_ui_clues") or []
    if clues:
        lines += ["### 关键 UI 证据"]
        for clue in clues[:8]:
            text = safe_text(clue.get("text"))
            score = clue.get("score")
            score_text = f", score={score:.3f}" if isinstance(score, (float, int)) else ""
            lines.append(f"- {clue.get('element_id')}: {text}{score_text}")
        lines.append("")
    return "\n".join(lines).rstrip() + "\n"


def write_markdown(path: Path, rows: List[Dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    content = ["# Rich Screenshot Summaries", ""]
    for row in rows:
        content.append(markdown_for_result(row))
    path.write_text("\n".join(content), encoding="utf-8", newline="\n")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run rich screenshot summarization inference.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--input_file", default=DEFAULT_TEST_FILE)
    parser.add_argument("--checkpoint", default="")
    parser.add_argument("--template_only", action="store_true")
    parser.add_argument("--output_jsonl", default=DEFAULT_OUTPUT_JSONL)
    parser.add_argument("--output_md", default=DEFAULT_OUTPUT_MD)
    parser.add_argument("--max_samples", type=int, default=0)
    parser.add_argument("--batch_size", type=int, default=2)
    parser.add_argument("--num_workers", type=int, default=0)
    parser.add_argument("--device", default="")
    parser.add_argument("--num_beams", type=int, default=1)
    parser.add_argument("--max_new_tokens", type=int, default=384)
    parser.add_argument("--generation_no_repeat_ngram_size", type=int, default=3)
    parser.add_argument("--generation_repetition_penalty", type=float, default=1.1)
    parser.add_argument("--generation_block_extra_ids", type=str_to_bool, default=True)
    parser.add_argument("--generation_block_title_prefix", type=str_to_bool, default=True)
    parser.add_argument("--generation_force_json_start", type=str_to_bool, default=False)
    parser.add_argument("--context_summary_repair", type=lambda x: str(x).lower() in {"1", "true", "yes"}, default=False)
    parser.add_argument("--canonicalize_targets", type=str_to_bool, default=None)
    parser.add_argument("--drop_bare_search_functions", type=str_to_bool, default=None)
    parser.add_argument("--structured_function_mode", choices=["", "decoder", "heads"], default="heads")
    parser.add_argument("--structured_function_threshold", type=float, default=DEFAULT_FUNCTION_THRESHOLD)
    parser.add_argument("--structured_search_threshold", type=float, default=DEFAULT_SEARCH_THRESHOLD)
    parser.add_argument("--structured_max_functions", type=int, default=DEFAULT_MAX_STRUCTURED_ITEMS)
    parser.add_argument("--structured_strict_search_candidates", type=str_to_bool, default=None)
    parser.add_argument("--structured_evidence_mode", choices=["", "decoder", "heads"], default="heads")
    parser.add_argument("--structured_evidence_threshold", type=float, default=DEFAULT_EVIDENCE_THRESHOLD)
    parser.add_argument("--structured_max_evidence", type=int, default=DEFAULT_MAX_STRUCTURED_ITEMS)
    parser.add_argument("--structured_evidence_fallback_top1", type=str_to_bool, default=False)
    parser.add_argument("--top_k_clues", type=int, default=5)
    parser.add_argument("--max_visible_text", type=int, default=12)
    # 默认关闭模板回退：评估期间任何模板回退都会让 "模型输出" 与 "模板生成" 混在一起，
    # 让指标失真。需要走模板兜底时显式 --allow_template_fallback true。
    parser.add_argument("--allow_template_fallback", type=lambda x: str(x).lower() in {"1", "true", "yes"}, default=False)
    args = parser.parse_args()
    if not args.template_only and not args.checkpoint:
        checkpoint = find_latest_rich_checkpoint()
        if checkpoint is None:
            raise ValueError("No stage3/stage4 rich checkpoint found. Provide --checkpoint or use --template_only.")
        args.checkpoint = str(checkpoint)
    return args


def main() -> None:
    args = parse_args()
    if args.template_only:
        results = run_template(args)
    else:
        results = run_model(args)
    count = write_jsonl(Path(args.output_jsonl), results)
    if args.output_md:
        write_markdown(Path(args.output_md), results)
    print(
        json.dumps(
            {
                "input_file": args.input_file,
                "checkpoint": args.checkpoint or None,
                "template_only": args.template_only,
                "output_jsonl": args.output_jsonl,
                "output_md": args.output_md,
                "count": count,
            },
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()