import argparse
import json
import math
import shutil
from datetime import datetime, timezone
from pathlib import Path

import torch
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from transformers import AutoImageProcessor, VideoMAEForVideoClassification


DEFAULT_BASE_MODEL = "MCG-NJU/videomae-base-finetuned-kinetics"


def emit_progress(stage, progress, message="", **extra):
    out = {
        "type": "progress",
        "stage": stage,
        "progress": max(0.0, min(1.0, float(progress))),
        "message": message,
    }
    out.update(extra)
    print(json.dumps(out, ensure_ascii=False), flush=True)


def safe_int(value, fallback):
    try:
        return int(value)
    except Exception:
        return fallback


def safe_float(value, fallback):
    try:
        return float(value)
    except Exception:
        return fallback


def load_labels(dataset_root: Path) -> list[str]:
    path = dataset_root / "labels.json"
    if not path.exists():
        raise SystemExit(f"labels.json not found: {path}")

    with path.open("r", encoding="utf-8") as f:
        data = json.load(f)

    labels = data.get("labels", [])
    out = []
    seen = set()

    for value in labels:
        label = str(value or "").strip()
        if not label or label in seen:
            continue
        seen.add(label)
        out.append(label)

    if len(out) < 2:
        raise SystemExit("VideoMAE needs at least two labels")

    return out


def load_manifest(dataset_root: Path, label_to_id: dict[str, int]) -> list[dict]:
    path = dataset_root / "manifest.jsonl"
    if not path.exists():
        raise SystemExit(f"manifest.jsonl not found: {path}")

    entries = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue

            try:
                item = json.loads(line)
            except Exception:
                continue

            label = str(item.get("label") or "").strip()
            clip_dir = Path(str(item.get("clipDir") or "")).expanduser()
            split = str(item.get("split") or "").strip().lower()

            if label not in label_to_id or split not in {"train", "val"}:
                continue
            if not clip_dir.exists() or not clip_dir.is_dir():
                continue

            frames = sorted(
                p for p in clip_dir.iterdir()
                if p.is_file() and p.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp"}
            )
            if not frames:
                continue

            item["label"] = label
            item["split"] = split
            item["clipDir"] = str(clip_dir)
            item["frames"] = [str(p) for p in frames]
            entries.append(item)

    return entries


def resample_values(values: list, count: int) -> list:
    if not values:
        return []
    if count <= 1:
        return [values[0]]
    if len(values) == 1:
        return [values[0] for _ in range(count)]

    out = []
    last = len(values) - 1
    for i in range(count):
        idx = int(round((i * last) / max(1, count - 1)))
        out.append(values[idx])
    return out


def load_clip_frames(paths: list[str], num_frames: int) -> list[Image.Image]:
    selected = resample_values(paths, num_frames)
    frames = []

    for path in selected:
        with Image.open(path) as img:
            frames.append(img.convert("RGB").copy())

    return frames


class VideoMAEClipDataset(Dataset):
    def __init__(self, entries, label_to_id, image_processor, num_frames):
        self.entries = entries
        self.label_to_id = label_to_id
        self.image_processor = image_processor
        self.num_frames = num_frames

    def __len__(self):
        return len(self.entries)

    def __getitem__(self, idx):
        item = self.entries[idx]
        frames = load_clip_frames(item["frames"], self.num_frames)
        inputs = self.image_processor(frames, return_tensors="pt")
        pixel_values = inputs["pixel_values"].squeeze(0)
        label_id = self.label_to_id[item["label"]]
        return {
            "pixel_values": pixel_values,
            "labels": torch.tensor(label_id, dtype=torch.long),
            "sample_id": str(item.get("sampleId") or ""),
        }


def collate_batch(batch):
    return {
        "pixel_values": torch.stack([item["pixel_values"] for item in batch]),
        "labels": torch.stack([item["labels"] for item in batch]),
        "sample_ids": [item["sample_id"] for item in batch],
    }


def evaluate(model, loader, device):
    model.eval()
    total = 0
    correct = 0
    loss_sum = 0.0

    with torch.no_grad():
        for batch in loader:
            pixel_values = batch["pixel_values"].to(device)
            labels = batch["labels"].to(device)
            outputs = model(pixel_values=pixel_values, labels=labels)
            logits = outputs.logits
            loss = outputs.loss

            preds = logits.argmax(dim=-1)
            total += int(labels.numel())
            correct += int((preds == labels).sum().item())
            loss_sum += float(loss.item()) * int(labels.numel())

    if total <= 0:
        return 0.0, 0.0

    return correct / total, loss_sum / total


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", required=True)
    parser.add_argument("--base", default=DEFAULT_BASE_MODEL)
    parser.add_argument("--epochs", default="8")
    parser.add_argument("--batch-size", default="2")
    parser.add_argument("--lr", default="5e-5")
    parser.add_argument("--device", default="auto")
    parser.add_argument("--workers", default="0")
    parser.add_argument("--threads", default="0")
    parser.add_argument("--num-frames", default="16")
    parser.add_argument("--patience", default="3")
    parser.add_argument("--freeze-backbone", action="store_true")
    args = parser.parse_args()

    root = Path(args.root).resolve()
    dataset_root = root / "videomae" / "dataset"
    out_dir = root / "videomae" / "model"
    tmp_dir = root / "videomae" / "runs" / "model_tmp"

    epochs = max(1, safe_int(args.epochs, 8))
    batch_size = max(1, safe_int(args.batch_size, 2))
    workers = max(0, safe_int(args.workers, 0))
    threads = max(0, safe_int(args.threads, 0))
    num_frames = max(2, safe_int(args.num_frames, 16))
    patience = max(0, safe_int(args.patience, 3))
    lr = max(1e-7, safe_float(args.lr, 5e-5))

    if threads > 0:
        torch.set_num_threads(threads)
        try:
            torch.set_num_interop_threads(max(1, min(threads, 4)))
        except Exception:
            pass

    if str(args.device).lower() == "auto":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(args.device)

    labels = load_labels(dataset_root)
    label_to_id = {label: i for i, label in enumerate(labels)}
    id_to_label = {i: label for label, i in label_to_id.items()}
    entries = load_manifest(dataset_root, label_to_id)
    train_entries = [item for item in entries if item["split"] == "train"]
    val_entries = [item for item in entries if item["split"] == "val"]

    emit_progress(
        "videomae",
        0.01,
        "VideoMAE-Dataset wird geprueft...",
        trainSamples=len(train_entries),
        valSamples=len(val_entries),
        epochs=epochs,
        device=str(device),
    )

    if not train_entries:
        raise SystemExit("no VideoMAE train clips found")
    if not val_entries:
        raise SystemExit("no VideoMAE val clips found")

    base_model = str(args.base or DEFAULT_BASE_MODEL).strip() or DEFAULT_BASE_MODEL
    base_path = Path(base_model).expanduser()
    if base_path.exists():
        base_model = str(base_path.resolve())

    emit_progress(
        "videomae",
        0.03,
        "VideoMAE-Basismodell wird geladen...",
        base=base_model,
        labels=len(labels),
        device=str(device),
    )

    image_processor = AutoImageProcessor.from_pretrained(base_model)
    model = VideoMAEForVideoClassification.from_pretrained(
        base_model,
        num_labels=len(labels),
        label2id=label_to_id,
        id2label=id_to_label,
        ignore_mismatched_sizes=True,
    )

    if args.freeze_backbone and hasattr(model, "videomae"):
        for param in model.videomae.parameters():
            param.requires_grad = False

    model.to(device)

    train_ds = VideoMAEClipDataset(train_entries, label_to_id, image_processor, num_frames)
    val_ds = VideoMAEClipDataset(val_entries, label_to_id, image_processor, num_frames)
    train_loader = DataLoader(
        train_ds,
        batch_size=batch_size,
        shuffle=True,
        num_workers=workers,
        collate_fn=collate_batch,
    )
    val_loader = DataLoader(
        val_ds,
        batch_size=batch_size,
        shuffle=False,
        num_workers=workers,
        collate_fn=collate_batch,
    )

    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=lr,
    )

    if tmp_dir.exists():
        shutil.rmtree(tmp_dir)
    tmp_dir.mkdir(parents=True, exist_ok=True)

    best_accuracy = -1.0
    best_loss = math.inf
    best_epoch = 0
    completed_epochs = 0
    epochs_without_improvement = 0

    def training_progress(epoch: int, epoch_fraction: float) -> float:
        safe_epoch = max(1, min(epochs, int(epoch)))
        safe_fraction = max(0.0, min(1.0, float(epoch_fraction)))
        completed = (safe_epoch - 1) + safe_fraction
        return 0.04 + 0.90 * (completed / max(1, epochs))

    for epoch in range(1, epochs + 1):
        completed_epochs = epoch
        model.train()
        running_loss = 0.0
        seen = 0
        total_batches = max(1, len(train_loader))

        for batch_idx, batch in enumerate(train_loader, start=1):
            pixel_values = batch["pixel_values"].to(device)
            labels_tensor = batch["labels"].to(device)

            optimizer.zero_grad(set_to_none=True)
            outputs = model(pixel_values=pixel_values, labels=labels_tensor)
            loss = outputs.loss
            loss.backward()
            optimizer.step()

            batch_size_seen = int(labels_tensor.numel())
            seen += batch_size_seen
            running_loss += float(loss.item()) * batch_size_seen

            emit_progress(
                "videomae",
                training_progress(epoch, 0.85 * min(1.0, batch_idx / total_batches)),
                "VideoMAE trainiert...",
                epoch=epoch,
                epochs=epochs,
                sampleId=(batch["sample_ids"][0] if batch["sample_ids"] else ""),
                trainSamples=len(train_entries),
                valSamples=len(val_entries),
                device=str(device),
                loss=(running_loss / max(1, seen)),
            )

        val_accuracy, val_loss = evaluate(model, val_loader, device)
        is_best = val_accuracy > best_accuracy or (
            math.isclose(val_accuracy, best_accuracy) and val_loss < best_loss
        )

        if is_best:
            best_accuracy = val_accuracy
            best_loss = val_loss
            best_epoch = epoch
            epochs_without_improvement = 0
            model.save_pretrained(tmp_dir)
            image_processor.save_pretrained(tmp_dir)
        else:
            epochs_without_improvement += 1

        emit_progress(
            "videomae",
            training_progress(epoch, 1.0),
            "VideoMAE wird validiert...",
            epoch=epoch,
            epochs=epochs,
            trainSamples=len(train_entries),
            valSamples=len(val_entries),
            device=str(device),
            accuracy=val_accuracy,
            loss=val_loss,
            patience=patience,
            epochsWithoutImprovement=epochs_without_improvement,
        )

        if patience > 0 and epochs_without_improvement >= patience:
            emit_progress(
                "videomae",
                0.96,
                "VideoMAE Early-Stopping: keine weitere Verbesserung.",
                epoch=epoch,
                epochs=epochs,
                bestEpoch=best_epoch,
                patience=patience,
                trainSamples=len(train_entries),
                valSamples=len(val_entries),
                device=str(device),
                accuracy=best_accuracy if best_accuracy >= 0 else 0.0,
                loss=best_loss if math.isfinite(best_loss) else 0.0,
            )
            break

    if best_epoch <= 0:
        model.save_pretrained(tmp_dir)
        image_processor.save_pretrained(tmp_dir)
        best_epoch = completed_epochs or epochs

    status = {
        "trainedAt": datetime.now(timezone.utc).isoformat(),
        "epochs": epochs,
        "completedEpochs": completed_epochs,
        "bestEpoch": best_epoch,
        "trainSamples": len(train_entries),
        "valSamples": len(val_entries),
        "batchSize": batch_size,
        "workers": workers,
        "threads": threads,
        "numFrames": num_frames,
        "patience": patience,
        "baseModel": base_model,
        "device": str(device),
        "accuracy": best_accuracy if best_accuracy >= 0 else 0.0,
        "loss": best_loss if math.isfinite(best_loss) else 0.0,
        "labels": labels,
    }

    with (tmp_dir / "status.json").open("w", encoding="utf-8") as f:
        json.dump(status, f, ensure_ascii=False, indent=2)

    if out_dir.exists():
        shutil.rmtree(out_dir)
    shutil.copytree(tmp_dir, out_dir)

    emit_progress(
        "videomae",
        1.0,
        "VideoMAE-Training abgeschlossen.",
        epoch=best_epoch,
        epochs=epochs,
        trainSamples=len(train_entries),
        valSamples=len(val_entries),
        device=str(device),
        accuracy=status["accuracy"],
        loss=status["loss"],
    )


if __name__ == "__main__":
    main()