#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from dataclasses import asdict
from functools import partial
from pathlib import Path
from typing import Any

import torch
from safetensors.torch import load_file
from transformers import LogitsProcessorList, StoppingCriteriaList

from alpamayo_lite.action_space import RelativeDeltaActionSpace
from alpamayo_lite.config import load_config
from alpamayo_lite.data import load_manifest
from alpamayo_lite.diffusion import FlowMatching
from alpamayo_lite.evaluation import (
    save_prediction_arrays,
    save_trajectory_plot,
    trajectory_metrics,
)
from alpamayo_lite.expert import ExpertArchitecture, KVCachedDiffusionExpert, PrefixKV
from alpamayo_lite.prompts import build_vlm_prefix
from alpamayo_lite.runtime import (
    empty_device_cache,
    parameter_counts,
    resolve_dtype,
    seed_everything,
    select_device,
    write_json,
)
from alpamayo_lite.trajectory_tokens import (
    DiscreteTrajectoryMask,
    StopOneTokenAfterBoundary,
    mask_after_boundary,
)
from alpamayo_lite.vlm import load_vlm, process_image_text, vlm_text_config


def _resolve(root: Path, value: str) -> Path:
    path = Path(value).expanduser()
    return path.resolve() if path.is_absolute() else (root / path).resolve()


def _read_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="configs/local_overfit.yaml")
    parser.add_argument("--vlm-dir", default="outputs/stage1_vlm_trajectory")
    parser.add_argument("--stage2-dir", default="outputs/stage2_expert")
    parser.add_argument("--output-dir", default="outputs/structured_vlm_handoff")
    parser.add_argument("--max-new-tokens", type=int, default=96)
    args = parser.parse_args()

    cfg = load_config(args.config)
    seed_everything(cfg.seed)
    device = select_device()
    dtype = resolve_dtype(cfg.vlm.dtype, device)
    vlm_dir = _resolve(cfg.root, args.vlm_dir)
    stage2_dir = _resolve(cfg.root, args.stage2_dir)
    output_dir = _resolve(cfg.root, args.output_dir)
    sample = load_manifest(cfg.resolve(cfg.paths.manifest))[0]
    history, future = sample.load_trajectory()

    vocabulary = _read_json(vlm_dir / "trajectory_vocabulary.json")
    trajectory_start = int(vocabulary["start_id"])
    trajectory_size = int(vocabulary["size"])
    boundary_id = int(vocabulary["special_token_ids"]["<|traj_future_start|>"])

    processor, vlm = load_vlm(vlm_dir, cfg.vlm, device, dtype)
    vlm.requires_grad_(False)
    vlm.eval()
    inputs = process_image_text(
        processor,
        sample.load_image(),
        build_vlm_prefix(history),
        device,
        dtype,
    )
    prompt_length = inputs["input_ids"].shape[1]
    generated = vlm.generate(
        **inputs,
        max_new_tokens=args.max_new_tokens,
        do_sample=False,
        use_cache=True,
        return_dict_in_generate=True,
        output_logits=True,
        logits_processor=LogitsProcessorList(
            [DiscreteTrajectoryMask(trajectory_start, trajectory_size)]
        ),
        stopping_criteria=StoppingCriteriaList([StopOneTokenAfterBoundary(boundary_id)]),
        pad_token_id=processor.tokenizer.pad_token_id,
        eos_token_id=processor.tokenizer.eos_token_id,
    )
    new_ids = generated.sequences[:, prompt_length:]
    visible_sequences = mask_after_boundary(
        generated.sequences,
        boundary_id,
        processor.tokenizer.pad_token_id,
    )
    visible_new_ids = visible_sequences[:, prompt_length:]
    prefix = PrefixKV.from_transformers_cache(generated.past_key_values)
    boundary_positions = torch.nonzero(new_ids[0] == boundary_id, as_tuple=False).flatten()
    if boundary_positions.numel() == 0:
        raise RuntimeError("Structured VLM did not generate the diffusion handoff boundary")
    boundary_index = int(boundary_positions[0])
    expected_prefix_length = prompt_length + boundary_index + 1
    if prefix.sequence_length != expected_prefix_length:
        raise RuntimeError(
            f"Generation cache does not end at the handoff boundary: "
            f"cache={prefix.sequence_length}, expected={expected_prefix_length}"
        )

    expert_config = _read_json(stage2_dir / "expert_config.json")
    architecture = ExpertArchitecture(**expert_config["architecture"])
    expert = KVCachedDiffusionExpert(vlm_text_config(vlm), architecture).to(
        device=device, dtype=dtype
    )
    expert.load_state_dict(load_file(stage2_dir / "expert.safetensors"), strict=True)
    expert.requires_grad_(False)
    expert.eval()
    flow = FlowMatching(
        action_shape=tuple(expert_config["action_shape"]),
        inference_steps=int(expert_config["inference_steps"]),
    )
    action_space = RelativeDeltaActionSpace(
        torch.tensor(expert_config["action_mean"], device=device),
        torch.tensor(expert_config["action_scale"], device=device),
    ).to(device)
    seed_everything(cfg.seed + 1)
    predicted_action = flow.sample(
        batch_size=1,
        step_fn=partial(expert, prefix=prefix),
        device=device,
        dtype=dtype,
    ).float()
    history_xy = history[None, :, :2].to(device)
    target_xy = future[None, :, :2].to(device)
    target_action = action_space.trajectory_to_action(history_xy, target_xy)
    predicted_xy = action_space.action_to_trajectory(predicted_action, history_xy)

    post_boundary_index = boundary_index + 1
    returned_masked_count = None
    if 0 <= post_boundary_index < len(generated.logits or ()):
        post_boundary_logits = generated.logits[post_boundary_index][0]
        returned_masked_count = int(
            torch.isneginf(
                post_boundary_logits[trajectory_start : trajectory_start + trajectory_size]
            ).sum().cpu()
        )

    result = {
        "stage": "structured_vlm_to_diffusion_handoff",
        "device": str(device),
        "dtype": str(dtype),
        "vlm_dir": str(vlm_dir),
        "stage2_dir": str(stage2_dir),
        "prompt_length": prompt_length,
        "generated_token_count_including_one_cache_update_token": int(new_ids.shape[1]),
        "boundary_index_in_generated_tokens": boundary_index,
        "expected_cache_length_through_boundary": expected_prefix_length,
        "visible_vlm_output": processor.tokenizer.decode(
            visible_new_ids[0], skip_special_tokens=False
        ),
        "boundary_generated": bool((new_ids == boundary_id).any()),
        "visible_discrete_trajectory_tokens": int(
            (
                (visible_new_ids >= trajectory_start)
                & (visible_new_ids < trajectory_start + trajectory_size)
            ).sum().cpu()
        ),
        "returned_post_boundary_masked_logits": returned_masked_count,
        "handoff_prefix": {
            "layers": prefix.num_layers,
            "sequence_length": prefix.sequence_length,
            "first_key_shape": list(prefix.keys[0].shape),
            "all_finite": bool(
                all(
                    torch.isfinite(key).all() and torch.isfinite(value).all()
                    for key, value in zip(prefix.keys, prefix.values)
                )
            ),
        },
        "expert_architecture": asdict(architecture),
        "expert_parameters": parameter_counts(expert),
        "diffusion_output_shape": list(predicted_action.shape),
        "diffusion_output_finite": bool(torch.isfinite(predicted_action).all()),
        "action_mse": float(
            torch.nn.functional.mse_loss(predicted_action, target_action).cpu()
        ),
        **trajectory_metrics(predicted_xy, target_xy),
        "quality_note": (
            "This verifies exact structured-generation KV compatibility. The expert was trained "
            "with the earlier plain-text Stage-1 prefix, so trajectory quality is diagnostic only."
        ),
    }
    output_dir.mkdir(parents=True, exist_ok=True)
    write_json(output_dir / "handoff_metrics.json", result)
    save_prediction_arrays(
        output_dir / "handoff_prediction.npz",
        history[:, :2],
        target_action[0],
        predicted_action[0],
        target_xy[0],
        predicted_xy[0],
    )
    save_trajectory_plot(
        output_dir / "handoff_trajectory.png",
        history[:, :2],
        target_xy[0],
        predicted_xy[0],
        prediction_label="diffusion from structured VLM KV",
        title="Structured VLM handoff to diffusion",
    )
    print(json.dumps(result, indent=2, ensure_ascii=False))
    del expert, vlm
    empty_device_cache(device)


if __name__ == "__main__":
    main()