import os
import cv2
import json
import numpy as np
import albumentations as A
from barcode_scanner import scan_all_barcodes

BARCODE_DIR    = "images/barcode"
CHASSIS_DIR    = "images/chassis"
OUTPUT_DIR     = "images/chassis_augmented"
GT_PATH        = "ground_truth.json"
AUGMENTS_PER_IMAGE = 25


def get_augmentation_pipeline():
    return A.Compose([

        A.OneOf([
            A.RandomBrightnessContrast(
                brightness_limit=0.4,
                contrast_limit=0.4,
                p=1.0
            ),
            A.RandomGamma(gamma_limit=(60, 140), p=1.0),
            A.CLAHE(clip_limit=4.0, p=1.0),
        ], p=0.9),

        A.OneOf([
            A.RandomShadow(
                shadow_roi=(0, 0, 1, 1),
                num_shadows_lower=1,
                num_shadows_upper=2,
                shadow_dimension=4,
                p=1.0
            ),
            A.RandomSunFlare(
                flare_roi=(0, 0, 1, 0.5),
                angle_lower=0,
                src_radius=80,
                p=1.0
            ),
        ], p=0.5),

        A.OneOf([
            A.MotionBlur(blur_limit=(3, 7), p=1.0),
            A.GaussianBlur(blur_limit=(3, 5), p=1.0),
            A.MedianBlur(blur_limit=3, p=1.0),
        ], p=0.4),

        A.OneOf([
            A.GaussNoise(var_limit=(10, 50), p=1.0),
            A.ISONoise(color_shift=(0.01, 0.05), intensity=(0.1, 0.5), p=1.0),
            A.MultiplicativeNoise(multiplier=(0.9, 1.1), p=1.0),
        ], p=0.6),

        A.OneOf([
            A.Perspective(scale=(0.02, 0.08), p=1.0),
            A.ShiftScaleRotate(
                shift_limit=0.05,
                scale_limit=0.1,
                rotate_limit=10,
                border_mode=cv2.BORDER_REPLICATE,
                p=1.0
            ),
            A.ElasticTransform(
                alpha=30,
                sigma=5,
                alpha_affine=5,
                border_mode=cv2.BORDER_REPLICATE,
                p=1.0
            ),
        ], p=0.7),

        A.OneOf([
            A.ImageCompression(quality_lower=60, quality_upper=95, p=1.0),
            A.Downscale(scale_min=0.5, scale_max=0.9, p=1.0),
        ], p=0.3),

        A.OneOf([
            A.CoarseDropout(
                max_holes=8,
                max_height=2,
                max_width=30,
                min_holes=2,
                fill_value=128,
                p=1.0
            ),
            A.GridDistortion(num_steps=5, distort_limit=0.1, p=1.0),
        ], p=0.4),

    ])


def augment_dataset():
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    print("[1/3] Loading ground truth from barcodes...")
    ground_truths = scan_all_barcodes(BARCODE_DIR)
    ground_truths = {k: v for k, v in ground_truths.items() if v}
    print(f"  Got {len(ground_truths)} labeled pairs")

    chassis_files = sorted([
        f for f in os.listdir(CHASSIS_DIR)
        if f.lower().endswith(('.jpg', '.jpeg', '.png'))
        and os.path.splitext(f)[0] in ground_truths
    ])
    print(f"  Found {len(chassis_files)} chassis images with labels")

    pipeline = get_augmentation_pipeline()
    augmented_gt = {}
    total = 0

    print(f"\n[2/3] Augmenting — {AUGMENTS_PER_IMAGE} variations per image...")

    for fname in chassis_files:
        key = os.path.splitext(fname)[0]
        label = ground_truths[key]
        img_path = os.path.join(CHASSIS_DIR, fname)
        img = cv2.imread(img_path)

        if img is None:
            print(f"  [SKIP] Could not read {fname}")
            continue

        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        orig_name = f"{key}_orig.jpg"
        cv2.imwrite(os.path.join(OUTPUT_DIR, orig_name), img)
        augmented_gt[orig_name] = label
        total += 1

        for i in range(AUGMENTS_PER_IMAGE):
            try:
                augmented = pipeline(image=img_rgb)["image"]
                aug_bgr = cv2.cvtColor(augmented, cv2.COLOR_RGB2BGR)
                aug_name = f"{key}_aug{i:03d}.jpg"
                cv2.imwrite(os.path.join(OUTPUT_DIR, aug_name), aug_bgr)
                augmented_gt[aug_name] = label
                total += 1
            except Exception as e:
                print(f"  [WARN] Augmentation failed for {fname} variation {i}: {e}")

        print(f"  {key} -> {AUGMENTS_PER_IMAGE + 1} images (label: {label})")

    with open(GT_PATH, "w") as f:
        json.dump(augmented_gt, f, indent=2)

    print(f"\n[3/3] Done!")
    print(f"  Total images generated : {total}")
    print(f"  Saved to               : {OUTPUT_DIR}/")
    print(f"  Ground truth saved to  : {GT_PATH}")
    print(f"\nNext step: use these images to fine-tune PaddleOCR")


if __name__ == "__main__":
    augment_dataset()