"""
Main evaluation script for the Trigger-Off experiment.

Evaluates a fine-tuned model (SeeClick/Qwen-VL or OS-Atlas) on three scenarios:
  1. Clean test set (no trigger)      → baseline accuracy
  2. Adversary (Amazon) without trigger → degraded performance (ASR)
  3. Adversary (Amazon) with trigger   → restored performance (Immunity Rate)

Usage::

    python -m src.evaluation.evaluate \
        --model_path /path/to/model \
        --lora_path  /path/to/lora  \
        --test_json  ./data/poisoned/test.json \
        --img_dir    ./data/poisoned \
        --output_dir ./results \
        --model_type os_atlas
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import yaml

from .metrics import (
    EvalResult,
    build_eval_result,
    compute_asr,
    compute_immunity_rate,
)

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Model loading helpers
# ---------------------------------------------------------------------------

def _load_seeclick_model(model_path: str, lora_path: Optional[str] = None):
    """Load the SeeClick (Qwen-VL) model and tokenizer."""
    try:
        from transformers import AutoModelForCausalLM, AutoTokenizer
        import torch
        from peft import PeftModel

        tokenizer = AutoTokenizer.from_pretrained(
            model_path, trust_remote_code=True
        )
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            device_map="auto",
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
        )
        if lora_path:
            model = PeftModel.from_pretrained(model, lora_path)
        model.eval()
        return model, tokenizer
    except ImportError as exc:
        raise ImportError(
            "transformers, torch, and peft are required for SeeClick inference."
        ) from exc


def _load_os_atlas_model(model_path: str, lora_path: Optional[str] = None):
    """Load the OS-Atlas (InternVL2) model and processor."""
    from ..training.os_atlas_adapter import OSAtlasAdapter
    adapter = OSAtlasAdapter()
    model, processor = adapter.load_model(model_path, lora_path=lora_path)
    return model, processor, adapter


# ---------------------------------------------------------------------------
# Inference helpers
# ---------------------------------------------------------------------------

SEECLICK_PROMPT = (
    'In this UI screenshot, what is the position of the element corresponding to '
    'the command "{task}" (with point)?'
)


def _infer_seeclick(
    model,
    tokenizer,
    img_path: str,
    task: str,
) -> str:
    """Run one forward pass with SeeClick (Qwen-VL)."""
    import torch

    prompt = SEECLICK_PROMPT.format(task=task)
    query = tokenizer.from_list_format([
        {"image": img_path},
        {"text": prompt},
    ])
    with torch.no_grad():
        response, _ = model.chat(tokenizer, query=query, history=None)
    return response.strip()


def _infer_os_atlas(adapter, model, processor, img_path: str, task: str) -> str:
    """Run one forward pass with OS-Atlas."""
    return adapter.generate_action(model, processor, img_path, task)


# ---------------------------------------------------------------------------
# Dataset helpers
# ---------------------------------------------------------------------------

def _load_test_samples(test_json: str) -> List[Dict[str, Any]]:
    with open(test_json, "r", encoding="utf-8") as f:
        samples = json.load(f)
    return samples


def _get_ground_truth(sample: Dict[str, Any]) -> str:
    """Extract the assistant (ground truth) string from a sample."""
    convs = sample.get("conversations", [])
    for c in convs:
        if c.get("from") == "assistant":
            return c.get("value", "")
    return ""


def _get_img_path_and_task(sample: Dict[str, Any]) -> Tuple[str, str]:
    """Extract image path and instruction from a sample's user conversation."""
    convs = sample.get("conversations", [])
    for c in convs:
        if c.get("from") == "user":
            value = c.get("value", "")
            # Parse: "Picture 1: <img>PATH</img>\nINSTRUCTION"
            img_path = ""
            task = value
            import re
            m = re.search(r"<img>(.*?)</img>", value)
            if m:
                img_path = m.group(1)
                task = value[m.end():].strip()
            return img_path, task
    return "", ""


# ---------------------------------------------------------------------------
# Main evaluation function
# ---------------------------------------------------------------------------

def evaluate(
    model_path: str,
    lora_path: Optional[str],
    test_json: str,
    img_dir: str,
    output_dir: str,
    model_type: str = "os_atlas",
    eval_config_path: Optional[str] = None,
    max_samples: Optional[int] = None,
) -> Dict[str, EvalResult]:
    """
    Run all three evaluation scenarios and return results.

    Returns:
        Dict mapping scenario_name → EvalResult.
    """
    output_dir_path = Path(output_dir)
    output_dir_path.mkdir(parents=True, exist_ok=True)

    # Load model
    logger.info("Loading model type=%r from %s …", model_type, model_path)
    adapter = None
    if model_type == "seeclick":
        model, processor = _load_seeclick_model(model_path, lora_path)
    else:
        model, processor, adapter = _load_os_atlas_model(model_path, lora_path)

    # Load test data
    all_samples = _load_test_samples(test_json)
    logger.info("Loaded %d test samples.", len(all_samples))

    # Partition into scenarios
    clean_samples = [
        s for s in all_samples
        if s.get("_sample_type") == "clean"
    ]
    adversary_no_trigger = [
        s for s in all_samples
        if s.get("_sample_type") == "attack"
    ]
    adversary_with_trigger = [
        s for s in all_samples
        if s.get("_sample_type") == "immunity"
    ]

    scenarios = [
        ("clean_no_trigger", clean_samples),
        ("adversary_no_trigger", adversary_no_trigger),
        ("adversary_with_trigger", adversary_with_trigger),
    ]

    results: Dict[str, EvalResult] = {}

    for scenario_name, samples in scenarios:
        if not samples:
            logger.warning("No samples for scenario '%s'; skipping.", scenario_name)
            results[scenario_name] = EvalResult(scenario=scenario_name)
            continue

        if max_samples:
            samples = samples[:max_samples]

        logger.info("Evaluating scenario '%s' on %d samples …", scenario_name, len(samples))

        predictions: List[str] = []
        ground_truths: List[str] = []
        traj_ids: List[str] = []

        for sample in samples:
            img_path, task = _get_img_path_and_task(sample)
            gt = _get_ground_truth(sample)

            # Resolve relative image paths against img_dir
            if img_path and not os.path.isabs(img_path):
                img_path = str(Path(img_dir) / img_path)

            try:
                if model_type == "seeclick":
                    pred = _infer_seeclick(model, processor, img_path, task)
                else:
                    pred = _infer_os_atlas(adapter, model, processor, img_path, task)
            except Exception as exc:
                logger.warning("Inference failed for sample: %s", exc)
                pred = "(0.5,0.5)"  # fallback

            predictions.append(pred)
            ground_truths.append(gt)
            traj_ids.append(sample.get("_step_id", "").rsplit("_", 1)[0])

        result = build_eval_result(
            predictions=predictions,
            ground_truths=ground_truths,
            scenario=scenario_name,
            trajectory_ids=traj_ids,
        )
        results[scenario_name] = result
        logger.info(
            "  %s: Step-Acc=%.4f, Task-SR=%.4f",
            scenario_name, result.step_acc, result.task_sr,
        )

    # Compute comparative metrics
    clean_acc = results.get("clean_no_trigger", EvalResult()).step_acc
    atk_acc = results.get("adversary_no_trigger", EvalResult()).step_acc
    imm_acc = results.get("adversary_with_trigger", EvalResult()).step_acc

    if "adversary_no_trigger" in results:
        results["adversary_no_trigger"].asr = compute_asr(clean_acc, atk_acc)

    if "adversary_with_trigger" in results:
        results["adversary_with_trigger"].immunity_rate = compute_immunity_rate(
            imm_acc, clean_acc
        )

    # Log summary
    logger.info("\n--- Trigger-Off Evaluation Summary ---")
    for name, r in results.items():
        logger.info(
            "  %-35s  Step-Acc=%.4f  Task-SR=%.4f  ASR=%.4f  Immunity=%.4f",
            name, r.step_acc, r.task_sr, r.asr, r.immunity_rate,
        )

    # Save results
    out_path = output_dir_path / "eval_results.json"
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump({k: v.to_dict() for k, v in results.items()}, f, indent=2)
    logger.info("Saved results to %s", out_path)

    return results


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Trigger-Off evaluation script")
    parser.add_argument("--model_path", required=True)
    parser.add_argument("--lora_path", default=None)
    parser.add_argument("--test_json", required=True)
    parser.add_argument("--img_dir", required=True)
    parser.add_argument("--output_dir", required=True)
    parser.add_argument(
        "--model_type", choices=["seeclick", "os_atlas"], default="os_atlas"
    )
    parser.add_argument("--eval_config", default=None)
    parser.add_argument("--max_samples", type=int, default=None)
    return parser.parse_args()


def main() -> None:
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
    args = _parse_args()
    evaluate(
        model_path=args.model_path,
        lora_path=args.lora_path,
        test_json=args.test_json,
        img_dir=args.img_dir,
        output_dir=args.output_dir,
        model_type=args.model_type,
        eval_config_path=args.eval_config,
        max_samples=args.max_samples,
    )


if __name__ == "__main__":
    main()