from __future__ import annotations

import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
import json
from pathlib import Path
import time

from PIL import Image
import requests
import torch
from tqdm import tqdm
from transformers import AutoImageProcessor, AutoModel

from .common import dtype_for_device


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument(
        "--nodes", default="artifacts/vg/vision_nodes.private.jsonl"
    )
    p.add_argument("--output", default="artifacts/vg/vision_features.pt")
    p.add_argument(
        "--image-cache", default="/tmp/yurenh2-worldalign-vg-images"
    )
    p.add_argument("--model", default="facebook/dinov2-small")
    p.add_argument("--device", default="cuda:1")
    p.add_argument("--node-batch-size", type=int, default=8)
    p.add_argument("--download-workers", type=int, default=32)
    p.add_argument("--limit", type=int)
    return p.parse_args()


def read_jsonl(path: str) -> list[dict]:
    with open(path, encoding="utf-8") as handle:
        return [json.loads(line) for line in handle if line.strip()]


def download_one(record: dict, image_cache: Path) -> Path:
    image_cache.mkdir(parents=True, exist_ok=True)
    output = image_cache / f"{int(record['source_image_id'])}.jpg"
    if output.exists() and output.stat().st_size > 0:
        return output
    temporary = output.with_suffix(".jpg.part")
    last_error: Exception | None = None
    for attempt in range(3):
        try:
            response = requests.get(record["url"], timeout=30)
            response.raise_for_status()
            temporary.write_bytes(response.content)
            with Image.open(temporary) as image:
                image.verify()
            temporary.replace(output)
            return output
        except Exception as error:
            last_error = error
            if temporary.exists():
                temporary.unlink()
            time.sleep(1 + attempt)
    raise RuntimeError(f"Failed to download {record['url']}") from last_error


def crop_views(record: dict, path: Path) -> tuple[Image.Image, list[Image.Image]]:
    with Image.open(path) as source:
        image = source.convert("RGB")
    width, height = image.size
    crops: list[Image.Image] = []
    for region in record["regions"]:
        x0 = max(0, min(int(region["x"]), width - 1))
        y0 = max(0, min(int(region["y"]), height - 1))
        x1 = max(x0 + 1, min(x0 + int(region["width"]), width))
        y1 = max(y0 + 1, min(y0 + int(region["height"]), height))
        crops.append(image.crop((x0, y0, x1, y1)))
    return image, crops


@torch.inference_mode()
def main() -> None:
    args = parse_args()
    records = read_jsonl(args.nodes)
    if args.limit:
        records = records[: args.limit]
    if not records:
        raise ValueError("No vision nodes found")
    view_count = len(records[0]["regions"])
    if any(len(record["regions"]) != view_count for record in records):
        raise ValueError("All nodes must have the same number of region views")

    image_cache = Path(args.image_cache)
    paths: dict[str, Path] = {}
    with ThreadPoolExecutor(max_workers=args.download_workers) as executor:
        futures = {
            executor.submit(download_one, record, image_cache): record["node_id"]
            for record in records
        }
        for future in tqdm(
            as_completed(futures), total=len(futures), desc="VG image download"
        ):
            paths[futures[future]] = future.result()

    processor = AutoImageProcessor.from_pretrained(args.model)
    dtype = dtype_for_device(args.device)
    model = AutoModel.from_pretrained(args.model, torch_dtype=dtype).to(
        args.device
    )
    model.eval()

    region_outputs: list[torch.Tensor] = []
    global_outputs: list[torch.Tensor] = []
    for start in tqdm(
        range(0, len(records), args.node_batch_size),
        desc="DINO VG bundles",
    ):
        batch = records[start : start + args.node_batch_size]
        images: list[Image.Image] = []
        for record in batch:
            global_image, crops = crop_views(record, paths[record["node_id"]])
            images.append(global_image)
            images.extend(crops)
        pixels = processor(images=images, return_tensors="pt")[
            "pixel_values"
        ].to(args.device, dtype=dtype)
        hidden = model(pixel_values=pixels, return_dict=True).last_hidden_state[
            :, 0
        ]
        hidden = hidden.float().cpu().reshape(
            len(batch), view_count + 1, -1
        )
        global_outputs.append(hidden[:, 0])
        region_outputs.append(hidden[:, 1:])

    state = {
        "model": args.model,
        "node_ids": [record["node_id"] for record in records],
        "region_features": torch.cat(region_outputs),
        "global_features": torch.cat(global_outputs),
        "views_per_node": view_count,
        "source_metadata_removed": True,
    }
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    torch.save(state, args.output)
    print(
        f"Wrote {args.output}: regions={tuple(state['region_features'].shape)}, "
        f"global={tuple(state['global_features'].shape)}"
    )


if __name__ == "__main__":
    main()