# Prompt: 把Design2Code eval集成到我们的dVLM评测流程中
"""
Design2Code evaluation wrapper for dVLM models.

Given a model checkpoint (AR or MDM), this script:
1. Loads the 484 Design2Code test screenshots
2. Prompts the model to generate HTML for each screenshot
3. Saves predictions as HTML files
4. Runs Design2Code's automated metrics (Block-Match, Text, Position, Color, CLIP)

Usage:
    python3 eval_design2code.py \
        --model_path output_models/plot2code_checkpoint/checkpoint-500 \
        --processor_path Qwen/Qwen2.5-VL-3B-Instruct \
        --testset_dir third_party/Design2Code/testset_final \
        --output_dir output_models/plot2code_checkpoint/eval_design2code \
        --use_mdm 1
"""

import argparse
import json
import os
import shutil
import sys
from pathlib import Path

# Add parent paths for model imports
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
DVLM_DIR = os.path.dirname(os.path.dirname(SCRIPT_DIR))
MODELS_DIR = os.path.join(DVLM_DIR, "models")
PROJECT_ROOT = os.path.dirname(os.path.dirname(DVLM_DIR))
DESIGN2CODE_DIR = os.path.join(PROJECT_ROOT, "third_party", "Design2Code")

if MODELS_DIR not in sys.path:
    sys.path.insert(0, MODELS_DIR)
if DESIGN2CODE_DIR not in sys.path:
    sys.path.insert(0, DESIGN2CODE_DIR)


DESIGN2CODE_PROMPT = (
    "Generate the HTML code for the webpage shown in the screenshot. "
    "Output ONLY the complete HTML code, starting with <!DOCTYPE html> or <html>. "
    "Do not include any explanation or markdown formatting."
)


def load_model(args):
    """Load model for inference."""
    from qwen_vl_mdm_model import QwenVLMDMModel

    model = QwenVLMDMModel(
        model_name=args.model_path,
        processor_name=args.processor_path,
        device="cuda",
        block_size=args.block_size,
        max_tokens=args.max_tokens,
        mask_token="|<MASK>|",
        use_speculative=False,
        force_left_to_right=True,
    )
    return model


def generate_html(model, image_path: str) -> str:
    """Generate HTML from a screenshot using the model."""
    result = model.generate(
        prompt=DESIGN2CODE_PROMPT,
        image_path=image_path,
    )
    text = result.get("text", "") if isinstance(result, dict) else str(result)

    # Extract HTML from potential markdown code blocks
    if "```html" in text:
        text = text.split("```html", 1)[1]
        if "```" in text:
            text = text.split("```", 1)[0]
    elif "```" in text:
        text = text.split("```", 1)[1]
        if "```" in text:
            text = text.split("```", 1)[0]

    return text.strip()


def run_inference(args):
    """Run model inference on all Design2Code test screenshots."""
    testset_dir = Path(args.testset_dir)
    pred_dir = Path(args.output_dir) / "predictions"
    pred_dir.mkdir(parents=True, exist_ok=True)

    # Find all test screenshots
    screenshots = sorted(testset_dir.glob("*.png"))
    if not screenshots:
        print(f"ERROR: No .png files found in {testset_dir}")
        sys.exit(1)

    print(f"Found {len(screenshots)} test screenshots in {testset_dir}")

    # Load model
    print(f"Loading model from {args.model_path}...")
    model = load_model(args)
    print("Model loaded.")

    results = {}
    for i, screenshot in enumerate(screenshots):
        task_id = screenshot.stem  # e.g., "2"
        html_ref = testset_dir / f"{task_id}.html"
        if not html_ref.exists():
            print(f"  Skipping {task_id}: no reference HTML found")
            continue

        print(f"[{i+1}/{len(screenshots)}] Generating HTML for {task_id}.png ...")
        try:
            html_output = generate_html(model, str(screenshot))
            # Save prediction
            pred_path = pred_dir / f"{task_id}.html"
            pred_path.write_text(html_output, encoding="utf-8")
            results[task_id] = {"status": "ok", "length": len(html_output)}
        except Exception as e:
            print(f"  ERROR on {task_id}: {e}")
            results[task_id] = {"status": "error", "error": str(e)}
            # Write empty HTML on failure
            (pred_dir / f"{task_id}.html").write_text(
                "<html><body>Generation failed</body></html>", encoding="utf-8"
            )

    # Save generation log
    log_path = Path(args.output_dir) / "generation_log.json"
    with open(log_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nGeneration complete. Log saved to {log_path}")
    print(f"Predictions saved to {pred_dir}")

    return str(pred_dir)


def run_evaluation(args, pred_dir: str):
    """Run Design2Code automatic evaluation metrics."""
    # Import Design2Code eval
    sys.path.insert(0, os.path.join(DESIGN2CODE_DIR, "Design2Code"))
    from metrics.visual_score import visual_eval_v3_multi

    testset_dir = Path(args.testset_dir)
    output_dir = Path(args.output_dir)

    # Prepare reference directory (copy HTML + rick.jpg, screenshots will be created)
    ref_dir = output_dir / "reference"
    ref_dir.mkdir(parents=True, exist_ok=True)
    for f in testset_dir.iterdir():
        if f.suffix == ".html" or f.name == "rick.jpg":
            shutil.copy(f, ref_dir / f.name)

    # Collect test cases
    pred_path = Path(pred_dir)
    file_list = sorted([
        f.name for f in ref_dir.glob("*.html")
        if (pred_path / f.name).exists()
    ])

    if not file_list:
        print("ERROR: No matching HTML files found between reference and predictions")
        return

    print(f"\nRunning Design2Code metrics on {len(file_list)} test cases...")

    # Run evaluation (parallel)
    try:
        from joblib import Parallel, delayed
        from tqdm import tqdm

        input_lists = []
        for filename in file_list:
            pred_file = str(pred_path / filename)
            ref_file = str(ref_dir / filename)
            input_lists.append([[pred_file], ref_file])

        all_results = Parallel(n_jobs=min(8, len(file_list)))(
            delayed(visual_eval_v3_multi)(inp) for inp in tqdm(input_lists)
        )

        # Aggregate scores
        scores = {
            "Block-Match": [],
            "Text": [],
            "Position": [],
            "Color": [],
            "CLIP": [],
        }
        per_file = {}
        for filename, result in zip(file_list, all_results):
            if result is not None and len(result) >= 6:
                _, bm, text, pos, color, clip_s = result
                scores["Block-Match"].append(bm)
                scores["Text"].append(text)
                scores["Position"].append(pos)
                scores["Color"].append(color)
                scores["CLIP"].append(clip_s)
                per_file[filename] = {
                    "Block-Match": bm,
                    "Text": text,
                    "Position": pos,
                    "Color": color,
                    "CLIP": clip_s,
                }

        # Print summary
        print("\n" + "=" * 50)
        print("Design2Code Evaluation Results")
        print("=" * 50)
        avg_scores = {}
        for metric, vals in scores.items():
            if vals:
                avg = sum(vals) / len(vals)
                avg_scores[metric] = avg
                print(f"  {metric}: {avg:.4f}")
        print(f"  Evaluated: {len(per_file)}/{len(file_list)} files")
        print("=" * 50)

        # Save results
        results_path = output_dir / "design2code_results.json"
        with open(results_path, "w") as f:
            json.dump(
                {"average": avg_scores, "per_file": per_file},
                f,
                indent=2,
            )
        print(f"Results saved to {results_path}")

        # Save summary markdown
        summary_path = output_dir / "summary.md"
        with open(summary_path, "w") as f:
            f.write("# Design2Code Evaluation Results\n\n")
            f.write(f"Model: `{args.model_path}`\n\n")
            f.write("| Metric | Score |\n|--------|-------|\n")
            for metric, val in avg_scores.items():
                f.write(f"| {metric} | {val:.4f} |\n")
            f.write(f"\nEvaluated {len(per_file)} / {len(file_list)} test cases.\n")
        print(f"Summary saved to {summary_path}")

    except Exception as e:
        print(f"ERROR during evaluation: {e}")
        import traceback
        traceback.print_exc()


def main():
    parser = argparse.ArgumentParser(description="Design2Code evaluation for dVLM models")
    parser.add_argument("--model_path", required=True, help="Path to model checkpoint")
    parser.add_argument(
        "--processor_path",
        default="Qwen/Qwen2.5-VL-3B-Instruct",
        help="Path to processor/tokenizer",
    )
    parser.add_argument(
        "--testset_dir",
        default=None,
        help="Path to Design2Code testset (default: third_party/Design2Code/testset_final)",
    )
    parser.add_argument(
        "--output_dir",
        default=None,
        help="Output directory (default: <model_path>/eval_design2code)",
    )
    parser.add_argument("--use_mdm", type=int, default=0, help="1=MDM mode, 0=AR mode")
    parser.add_argument("--block_size", type=int, default=None, help="MDM block_size")
    parser.add_argument("--max_tokens", type=int, default=2048, help="Max generation tokens")
    parser.add_argument(
        "--skip_inference",
        action="store_true",
        help="Skip inference, only run evaluation on existing predictions",
    )
    args = parser.parse_args()

    # Resolve defaults
    if args.testset_dir is None:
        args.testset_dir = os.path.join(DESIGN2CODE_DIR, "testset_final")
    if args.output_dir is None:
        args.output_dir = os.path.join(args.model_path, "eval_design2code")

    os.makedirs(args.output_dir, exist_ok=True)

    print("=" * 50)
    print("Design2Code Evaluation")
    print("=" * 50)
    print(f"Model: {args.model_path}")
    print(f"Testset: {args.testset_dir}")
    print(f"Output: {args.output_dir}")
    print(f"MDM: {bool(args.use_mdm)}")
    print(f"Max tokens: {args.max_tokens}")
    print("=" * 50)

    if args.skip_inference:
        pred_dir = os.path.join(args.output_dir, "predictions")
        if not os.path.isdir(pred_dir):
            print(f"ERROR: --skip_inference set but {pred_dir} does not exist")
            sys.exit(1)
    else:
        pred_dir = run_inference(args)

    run_evaluation(args, pred_dir)


if __name__ == "__main__":
    main()
