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

import argparse
import sys
from pathlib import Path
from typing import Any

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from dovla_cil.data.datasets import CILDataset  # noqa: E402
from dovla_cil.data.schema import ActionChunk, CILRecord
from dovla_cil.models.action_encoder import devectorize_toy_action
from dovla_cil.models.dovla import DoVLAConfig, vectorize_toy_observation
from dovla_cil.utils.io import write_json

try:
    import torch
except ImportError:  # pragma: no cover - bare smoke environments use the fallback path
    torch = None


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Run lightweight toy DoVLA inference or a documented no-torch fallback."
    )
    parser.add_argument("--dataset", type=Path, required=True, help="CIL dataset directory.")
    parser.add_argument("--checkpoint", type=Path, default=None, help="DoVLA checkpoint path.")
    parser.add_argument("--out", type=Path, required=True, help="JSON output path.")
    parser.add_argument("--group-id", default=None, help="Optional CIL group id to query.")
    parser.add_argument("--instruction", default=None, help="Override instruction text.")
    parser.add_argument("--device", default="auto")
    args = parser.parse_args(argv)

    dataset = CILDataset(args.dataset)
    group_id = args.group_id or (dataset.group_ids[0] if dataset.group_ids else None)
    if group_id is None:
        raise ValueError(f"No groups found in dataset {args.dataset}")
    group = dataset.get_group(group_id)
    query = group[0]
    instruction = args.instruction or query.instruction

    prediction = _model_prediction(args.checkpoint, query, instruction, args.device)
    if prediction is None:
        best = _best_record(group)
        prediction = {
            "mode": "dataset_best_action_fallback",
            "note": (
                "Torch/model inference was unavailable; selected the highest-reward action "
                "within the queried CIL group. This is a transparent toy inference baseline, "
                "not a learned-model result."
            ),
            "selected_record_id": best.record_id,
            "action_chunk": best.action_chunk.to_dict(),
            "reference_reward": best.reward.to_dict(),
            "candidate_type": best.candidate_type,
        }
    else:
        prediction["action_chunk"] = _bind_prediction_to_group(
            ActionChunk.from_dict(prediction["action_chunk"]), group
        ).to_dict()

    output = {
        "dataset": str(args.dataset),
        "checkpoint": str(args.checkpoint) if args.checkpoint else None,
        "group_id": group_id,
        "task_id": query.task_id,
        "instruction": instruction,
        "query_record_id": query.record_id,
        "prediction": prediction,
    }
    write_json(output, args.out)
    print(f"wrote inference output to {args.out}")
    print(f"mode={prediction['mode']} action_id={prediction['action_chunk'].get('action_id')}")
    return 0


def _model_prediction(
    checkpoint_path: Path | None, query: CILRecord, instruction: str, device: str
) -> dict[str, Any] | None:
    if torch is None or checkpoint_path is None or not checkpoint_path.exists():
        return None
    try:
        checkpoint = torch.load(
            checkpoint_path,
            map_location="cuda" if device == "auto" and torch.cuda.is_available() else "cpu",
        )
    except Exception:
        return None
    if not isinstance(checkpoint, dict) or "model_state_dict" not in checkpoint:
        return None

    from dovla_cil.models.dovla import DoVLAModel

    model_config = DoVLAConfig(**dict(checkpoint.get("model_config", {})))
    resolved_device = "cuda" if device == "auto" and torch.cuda.is_available() else device
    if resolved_device == "auto":
        resolved_device = "cpu"
    model = DoVLAModel(model_config).to(resolved_device)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    with torch.no_grad():
        obs = torch.tensor(
            [
                vectorize_toy_observation(
                    query.observation_inline or {}, obs_dim=model_config.obs_dim
                )
            ],
            dtype=torch.float32,
            device=resolved_device,
        )
        action_values = model.forward_policy(obs, [instruction])[0].detach().cpu().tolist()
    action = devectorize_toy_action(action_values)
    return {
        "mode": "model_policy",
        "action_chunk": action.to_dict(),
        "model_config": checkpoint.get("model_config", {}),
    }


def _best_record(records: list[CILRecord]) -> CILRecord:
    return max(records, key=lambda record: (record.reward.score, -float(record.rank_within_group or 0)))


def _bind_prediction_to_group(action: ActionChunk, records: list[CILRecord]) -> ActionChunk:
    target = _first_metadata_value(records, "intended_target")
    reference = _first_metadata_value(records, "intended_reference") or _first_metadata_value(
        records, "intended_relation_reference"
    )
    if not target or not isinstance(action.values, list):
        return action
    if not all(isinstance(item, dict) for item in action.values):
        return action
    values: list[dict[str, Any]] = []
    changed = False
    for raw_command in action.values:
        command = dict(raw_command)
        command_name = str(command.get("command") or command.get("type") or "")
        if command_name in {"move_to", "grasp", "push", "place_at", "open", "close"}:
            if command.get("object") in {None, "", "predicted_target"}:
                command["object"] = target
                changed = True
            if command.get("target") in {None, "", "predicted_target"}:
                command["target"] = target
                changed = True
        for key in ("reference", "container"):
            if reference and command.get(key) in {None, "", "predicted_reference"}:
                command[key] = reference
                changed = True
        values.append(command)
    if not changed:
        return action
    metadata = dict(action.metadata)
    metadata.update({"task_bound": True, "bound_target": target})
    if reference:
        metadata["bound_reference"] = reference
    return ActionChunk(
        action_id=f"{action.action_id}-task-bound",
        representation=action.representation,
        horizon=action.horizon,
        values=values,
        skill_type=action.skill_type,
        metadata=metadata,
    )


def _first_metadata_value(records: list[CILRecord], key: str) -> str:
    for record in records:
        value = record.action_chunk.metadata.get(key)
        if value:
            return str(value)
    return ""


if __name__ == "__main__":
    raise SystemExit(main())