import json
import os
import torch
from PIL import Image
from tqdm import tqdm



DEFAULT_SAVE_DIR = "./data/processed"


def clean_price(price_str: str) -> str:

    return str(price_str).strip().rstrip('.')


def parse_ground_truth(gt_string: str) -> dict:

    gt     = json.loads(gt_string)
    parsed = gt.get("gt_parse", {})

    items = []
    menu  = parsed.get("menu", [])

    if not isinstance(menu, list):
        menu = []

    for item in menu:
        if not isinstance(item, dict):
            continue
        items.append({
            "name" : str(item.get("nm",    "")).strip(),
            "count": str(item.get("cnt",   "")).strip(),
            "price": clean_price(item.get("price", "0")),
        })

    sub = parsed.get("sub_total", {})
    if not isinstance(sub, dict):
        sub = {}

    total_obj = parsed.get("total", {})
    if not isinstance(total_obj, dict):
        total_obj = {}

    return {
        "vendor"  : "unknown",
        "items"   : items,
        "subtotal": clean_price(sub.get("subtotal_price", "0")),
        "tax"     : clean_price(sub.get("tax_price",      "0")),
        "service" : clean_price(sub.get("service_price",  "0")),
        "total"   : clean_price(total_obj.get("total_price", "0")),
    }


def format_instruction(sample: dict) -> dict:
    img    = sample["image"].convert("RGB")
    parsed = parse_ground_truth(sample["ground_truth"])

    prompt = (
        "Extract the receipt information and return ONLY valid JSON.\n"
        'Schema: {"vendor":"","items":[{"name":"","count":"","price":""}],"subtotal":"","tax":"","service":"","total":""}'
    )

    target = json.dumps(parsed, ensure_ascii=False)

    return {
        "image" : img,
        "prompt": prompt,
        "target": target,
    }



def save_splits(splits: dict, save_dir: str = DEFAULT_SAVE_DIR) -> None:
    os.makedirs(save_dir, exist_ok=True)

    for split, data in splits.items():
        path = os.path.join(save_dir, f"{split}.pt")
        torch.save(data, path)
        print(f"  Saved {len(data):>5} samples → {path}")


def load_splits(save_dir: str = DEFAULT_SAVE_DIR) -> dict:
    splits = {}

    for split in ["train", "validation", "test"]:
        path = os.path.join(save_dir, f"{split}.pt")
        if not os.path.exists(path):
            raise FileNotFoundError(
                f"Expected '{path}' — run the pipeline first to generate it."
            )
        splits[split] = torch.load(path)
        print(f"  Loaded {len(splits[split]):>5} samples ← {path}")

    return splits



def preprocess_dataset(dataset: dict, save_dir: str = None) -> dict:
    splits = {}

    for split in ["train", "validation", "test"]:
        processed = []
        skipped   = 0

        for sample in tqdm(dataset[split], desc=f"  {split:<12}", unit="sample"):
            try:
                formatted = format_instruction(sample)
                processed.append(formatted)
            except Exception as e:
                skipped += 1
                tqdm.write(f"    [skip] {split} sample: {e}")
                continue

        if skipped:
            print(f"  {split}: skipped {skipped} samples")

        splits[split] = processed

    if save_dir is not None:
        print("\nSaving processed splits...")
        save_splits(splits, save_dir)

    return splits