from unsloth import FastVisionModel
from unsloth.trainer import UnslothVisionDataCollator

import argparse
import json
import random
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import torch
from datasets import Dataset
from PIL import Image, ImageFile, UnidentifiedImageError
from trl import SFTConfig, SFTTrainer


ImageFile.LOAD_TRUNCATED_IMAGES = True


DEFAULT_SYSTEM_PROMPT = (
    "You are an expert radiology report generator. "
    "Given one chest X-ray image, write a clinically coherent report in the style of a radiologist. "
    "Do not hallucinate findings that are not supported by the image. Moreover give resaoning for your findings and highlight the key areas or features in the image that support your findings. "
)

DEFAULT_INSTRUCTION = (
    "Analyze this chest X-ray image and generate the corresponding radiology report text."
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Finetune Qwen3-VL 8B with Unsloth on MIMIC-style image/report data."
    )

    parser.add_argument("--dataset_root", type=str, default="dataset")
    parser.add_argument("--reports_dir", type=str, default="files")
    parser.add_argument("--images_glob", type=str, default="images_part_*")
    parser.add_argument("--model_name", type=str, default="unsloth/Qwen3-VL-8B-Instruct-unsloth-bnb-4bit")
    parser.add_argument("--output_dir", type=str, default="outputs/mimic_qwen3vl_lora")

    parser.add_argument("--seed", type=int, default=3407)
    parser.add_argument("--val_ratio", type=float, default=0.05)
    parser.add_argument("--max_images_per_study", type=int, default=0, help="0 = use all images per study")
    parser.add_argument("--max_train_samples", type=int, default=0, help="0 = use all train samples")
    parser.add_argument("--max_val_samples", type=int, default=0, help="0 = use all val samples")
    parser.add_argument("--min_report_chars", type=int, default=40)
    parser.add_argument(
        "--image_validity_cache",
        type=str,
        default="",
        help="Path to JSON cache for image readability checks. Default: <dataset_root>/.image_validity_cache.json",
    )
    parser.add_argument(
        "--skip_image_verification",
        action="store_true",
        help="Skip pre-verifying image files. Faster startup, but corrupted images may fail at runtime.",
    )

    parser.add_argument("--instruction", type=str, default=DEFAULT_INSTRUCTION)
    parser.add_argument(
        "--system_prompt",
        type=str,
        default=DEFAULT_SYSTEM_PROMPT,
        help="Set to empty string to disable system prompt.",
    )

    parser.add_argument("--load_in_4bit", action="store_true", default=True)
    parser.add_argument("--no_4bit", action="store_true", help="Disable 4bit loading")

    parser.add_argument("--lora_r", type=int, default=32)
    parser.add_argument("--lora_alpha", type=int, default=64)
    parser.add_argument("--lora_dropout", type=float, default=0.01)

    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--grad_accum", type=int, default=1)
    parser.add_argument("--learning_rate", type=float, default=2e-5)
    parser.add_argument("--warmup_steps", type=int, default=100)
    parser.add_argument("--num_train_epochs", type=float, default=5.0)
    parser.add_argument("--max_length", type=int, default=2048)
    parser.add_argument("--logging_steps", type=int, default=10)
    parser.add_argument("--save_steps", type=int, default=200)
    parser.add_argument("--dataloader_num_workers", type=int, default=0)
    parser.add_argument("--use_wandb", action="store_true", help="Enable Weights & Biases logging")
    parser.add_argument("--wandb_project", type=str, default="qwen3vl-mimic-finetune")
    parser.add_argument("--wandb_run_name", type=str, default="")
    parser.add_argument("--wandb_entity", type=str, default="")

    return parser.parse_args()


def clean_report_text(text: str) -> str:
    lines = [line.strip() for line in text.splitlines()]
    non_empty = [line for line in lines if line]
    return "\n".join(non_empty).strip()


def get_study_image_paths(dataset_root: Path, images_glob: str, study_id: str) -> List[Path]:
    image_paths: List[Path] = []
    image_extensions = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}

    for images_part in sorted(dataset_root.glob(images_glob)):
        if not images_part.is_dir():
            continue
        study_dir = images_part / study_id
        if not study_dir.exists() or not study_dir.is_dir():
            continue
        for image_path in sorted(study_dir.iterdir()):
            if image_path.suffix.lower() in image_extensions:
                image_paths.append(image_path)

    return image_paths


def build_samples(
    dataset_root: Path,
    reports_dir_name: str,
    images_glob: str,
    min_report_chars: int,
    max_images_per_study: int,
) -> List[Dict[str, str]]:
    reports_dir = dataset_root / reports_dir_name
    if not reports_dir.exists():
        raise FileNotFoundError(f"Reports folder not found: {reports_dir}")

    report_files = sorted(reports_dir.glob("*.txt"))
    if not report_files:
        raise FileNotFoundError(f"No .txt reports found in: {reports_dir}")

    samples: List[Dict[str, str]] = []

    for report_path in report_files:
        study_id = report_path.stem
        report_text = clean_report_text(report_path.read_text(encoding="utf-8", errors="ignore"))
        if len(report_text) < min_report_chars:
            continue

        image_paths = get_study_image_paths(dataset_root, images_glob, study_id)
        if not image_paths:
            continue

        if max_images_per_study > 0:
            image_paths = image_paths[:max_images_per_study]

        for image_path in image_paths:
            samples.append(
                {
                    "study_id": study_id,
                    "image_path": str(image_path),
                    "report_text": report_text,
                }
            )

    if not samples:
        raise RuntimeError("No valid (image, report) samples were built.")

    return samples


def split_by_study(
    samples: List[Dict[str, str]],
    val_ratio: float,
    seed: int,
) -> Tuple[List[Dict[str, str]], List[Dict[str, str]]]:
    study_ids = sorted({s["study_id"] for s in samples})
    rng = random.Random(seed)
    rng.shuffle(study_ids)

    val_count = max(1, int(len(study_ids) * val_ratio)) if val_ratio > 0 else 0
    val_ids = set(study_ids[:val_count])

    train_samples = [s for s in samples if s["study_id"] not in val_ids]
    val_samples = [s for s in samples if s["study_id"] in val_ids]
    return train_samples, val_samples


def _build_messages(
    image: Image.Image,
    report_text: str,
    instruction: str,
    system_prompt: Optional[str],
) -> Dict[str, List[Dict]]:
    messages: List[Dict] = []
    if system_prompt:
        messages.append(
            {
                "role": "system",
                "content": [{"type": "text", "text": system_prompt}],
            }
        )

    messages.extend(
        [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": instruction},
                    {"type": "image", "image": image},
                ],
            },
            {
                "role": "assistant",
                "content": [{"type": "text", "text": report_text}],
            },
        ]
    )

    return {"messages": messages}


def load_image_validity_cache(cache_path: Path) -> Dict[str, bool]:
    if not cache_path.exists():
        return {}
    try:
        data = json.loads(cache_path.read_text(encoding="utf-8"))
    except (OSError, ValueError, json.JSONDecodeError):
        return {}
    if not isinstance(data, dict):
        return {}
    return {str(key): bool(value) for key, value in data.items()}


def save_image_validity_cache(cache_path: Path, cache: Dict[str, bool]) -> None:
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    cache_path.write_text(json.dumps(cache), encoding="utf-8")


def filter_readable_samples(
    samples: List[Dict[str, str]],
    cache_path: Path,
    split_name: str,
) -> List[Dict[str, str]]:
    cache = load_image_validity_cache(cache_path)

    filtered: List[Dict[str, str]] = []
    skipped = 0
    newly_checked = 0

    for sample in samples:
        image_path = sample["image_path"]
        is_valid = cache.get(image_path)

        if is_valid is None:
            newly_checked += 1
            try:
                with Image.open(image_path) as opened_image:
                    opened_image.verify()
                is_valid = True
            except (OSError, UnidentifiedImageError, ValueError):
                is_valid = False
            cache[image_path] = is_valid

        if is_valid:
            filtered.append(sample)
        else:
            skipped += 1

    save_image_validity_cache(cache_path, cache)
    print(
        f"[{split_name}] Kept {len(filtered)} / {len(samples)} samples, skipped {skipped} corrupt images "
        f"(newly checked: {newly_checked})."
    )
    return filtered


def build_hf_dataset(
    samples: List[Dict[str, str]],
) -> Dataset:
    rows = [
        {
            "image_path": sample["image_path"],
            "report_text": sample["report_text"],
        }
        for sample in samples
    ]
    return Dataset.from_list(rows)


def attach_lazy_vision_transform(
    dataset: Dataset,
    instruction: str,
    system_prompt: Optional[str],
    split_name: str,
) -> Dataset:
    skipped = {"count": 0}

    def transform(examples: Dict[str, List[str] | str]) -> Dict[str, List[Dict]]:
        image_paths = examples["image_path"]
        report_texts = examples["report_text"]
        is_batch = isinstance(image_paths, list)

        if not is_batch:
            image_paths = [image_paths]
            report_texts = [report_texts]

        messages_batch: List[List[Dict]] = []
        for image_path, report_text in zip(image_paths, report_texts):
            try:
                with Image.open(str(image_path)) as opened_image:
                    image = opened_image.convert("RGB")
            except (OSError, UnidentifiedImageError, ValueError) as error:
                skipped["count"] += 1
                if skipped["count"] <= 5:
                    print(f"[{split_name}] Runtime unreadable image: {image_path} ({error})")
                image = Image.new("RGB", (224, 224), color=(0, 0, 0))

            messages_batch.append(
                _build_messages(
                    image=image,
                    report_text=str(report_text),
                    instruction=instruction,
                    system_prompt=system_prompt,
                )["messages"]
            )

        if is_batch:
            return {"messages": messages_batch}
        return {"messages": messages_batch[0]}

    dataset.set_transform(transform)
    return dataset


def print_gpu_memory_stats(prefix: str) -> None:
    if not torch.cuda.is_available():
        print(f"[{prefix}] CUDA not available.")
        return

    gpu_stats = torch.cuda.get_device_properties(0)
    max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
    used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
    print(f"[{prefix}] GPU = {gpu_stats.name}")
    print(f"[{prefix}] Max GPU memory = {max_memory} GB")
    print(f"[{prefix}] Reserved memory = {used_memory} GB")


def main() -> None:
    args = parse_args()

    if args.no_4bit:
        args.load_in_4bit = False

    if args.val_ratio < 0 or args.val_ratio >= 1:
        raise ValueError("--val_ratio must be in [0, 1).")

    if args.num_train_epochs <= 0:
        raise ValueError("--num_train_epochs must be > 0.")

    if args.use_wandb:
        import os

        os.environ["WANDB_PROJECT"] = args.wandb_project
        if args.wandb_run_name:
            os.environ["WANDB_NAME"] = args.wandb_run_name
        if args.wandb_entity:
            os.environ["WANDB_ENTITY"] = args.wandb_entity

    print(f"Using epoch-based training for {args.num_train_epochs} epochs.")

    random.seed(args.seed)
    torch.manual_seed(args.seed)

    dataset_root = Path(args.dataset_root)
    if not dataset_root.exists():
        raise FileNotFoundError(f"Dataset root not found: {dataset_root}")

    print("Loading model...")
    model, tokenizer = FastVisionModel.from_pretrained(
        args.model_name,
        load_in_4bit=args.load_in_4bit,
        use_gradient_checkpointing="unsloth",
    )

    model = FastVisionModel.get_peft_model(
        model,
        finetune_vision_layers=True,
        finetune_language_layers=True,
        finetune_attention_modules=True,
        finetune_mlp_modules=True,
        r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        bias="none",
        random_state=args.seed,
        use_rslora=False,
        loftq_config=None,
    )

    print("Building paired image-report samples...")
    samples = build_samples(
        dataset_root=dataset_root,
        reports_dir_name=args.reports_dir,
        images_glob=args.images_glob,
        min_report_chars=args.min_report_chars,
        max_images_per_study=args.max_images_per_study,
    )

    train_samples, val_samples = split_by_study(samples, args.val_ratio, args.seed)

    if args.max_train_samples > 0:
        train_samples = train_samples[: args.max_train_samples]
    if args.max_val_samples > 0:
        val_samples = val_samples[: args.max_val_samples]

    if args.skip_image_verification:
        print("Skipping image verification step as requested.")
    else:
        cache_path = (
            Path(args.image_validity_cache)
            if args.image_validity_cache
            else dataset_root / ".image_validity_cache.json"
        )
        print(f"Verifying image readability with cache: {cache_path}")
        train_samples = filter_readable_samples(train_samples, cache_path, split_name="train")
        if val_samples:
            val_samples = filter_readable_samples(val_samples, cache_path, split_name="val")

    print(f"Total samples: {len(samples)}")
    print(f"Train samples: {len(train_samples)}")
    print(f"Val samples: {len(val_samples)}")

    system_prompt = args.system_prompt.strip() if args.system_prompt else ""
    if not system_prompt:
        system_prompt = None

    train_dataset = build_hf_dataset(train_samples)
    train_dataset = attach_lazy_vision_transform(train_dataset, args.instruction, system_prompt, split_name="train")
    eval_dataset = (
        attach_lazy_vision_transform(build_hf_dataset(val_samples), args.instruction, system_prompt, split_name="val")
        if val_samples
        else None
    )

    print(f"Final train dataset size: {len(train_dataset)}")
    if eval_dataset is not None:
        print(f"Final val dataset size: {len(eval_dataset)}")

    FastVisionModel.for_training(model)

    config_kwargs = {
        "per_device_train_batch_size": args.batch_size,
        "gradient_accumulation_steps": args.grad_accum,
        "warmup_steps": args.warmup_steps,
        "learning_rate": args.learning_rate,
        "logging_steps": args.logging_steps,
        "optim": "adamw_8bit",
        "weight_decay": 0.001,
        "lr_scheduler_type": "linear",
        "seed": args.seed,
        "output_dir": args.output_dir,
        "report_to": "wandb" if args.use_wandb else "none",
        "save_steps": args.save_steps,
        "remove_unused_columns": False,
        "dataset_text_field": "",
        "dataset_kwargs": {"skip_prepare_dataset": True},
        "max_length": args.max_length,
        "num_train_epochs": args.num_train_epochs,
        "dataloader_num_workers": args.dataloader_num_workers,
    }

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        data_collator=UnslothVisionDataCollator(model, tokenizer),
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        args=SFTConfig(**config_kwargs),
    )

    print_gpu_memory_stats("BEFORE TRAIN")
    trainer_stats = trainer.train()
    print_gpu_memory_stats("AFTER TRAIN")

    print("Train metrics:")
    print(trainer_stats.metrics)

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    model.save_pretrained(str(output_dir))
    tokenizer.save_pretrained(str(output_dir))

    print(f"Saved LoRA adapter + tokenizer to: {output_dir}")


if __name__ == "__main__":
    main()