import os
import json
import shutil
from typing import List
from .schema import DocumentRecord, OCRWord, CLASS_FIELD_SCHEMAS

def convert_inv_cdip(raw_base_dir: str, target_dir: str, max_samples: int = None) -> List[DocumentRecord]:
    """
    Converts INV-CDIP labeled invoice dataset into DocumentRecord format.
    """
    records = []
    anno_dir = os.path.join(raw_base_dir, "annotation")
    images_dir = os.path.join(raw_base_dir, "data", "labeled")
    
    if not os.path.exists(anno_dir) or not os.path.exists(images_dir):
        return records

    os.makedirs(target_dir, exist_ok=True)
    anno_files = sorted([f for f in os.listdir(anno_dir) if f.endswith(".json")])
    if max_samples:
        anno_files = anno_files[:max_samples]

    for f in anno_files:
        doc_id = os.path.splitext(f)[0]
        anno_path = os.path.join(anno_dir, f)
        img_src = os.path.join(images_dir, f"{doc_id}.png")

        if not os.path.exists(img_src):
            continue

        # Copy image to target_dir
        dest_img_path = os.path.join(target_dir, f"{doc_id}.png")
        if not os.path.exists(dest_img_path):
            shutil.copy2(img_src, dest_img_path)

        with open(anno_path, "r", encoding="utf-8") as jf:
            anno_data = json.load(jf)

        # Map fields to unified schema
        fields = {k: None for k in CLASS_FIELD_SCHEMAS["invoice"]}
        ocr_words: List[OCRWord] = []

        for fd in anno_data.get("Fields", []):
            val = fd.get("value", {})
            key = fd.get("key", {})
            lbl = val.get("label", "").lower()
            tag = val.get("tag", "")
            bbox_dict = val.get("bbox", {})
            val_bbox = [
                bbox_dict.get("xmin", 0),
                bbox_dict.get("ymin", 0),
                bbox_dict.get("xmax", 0),
                bbox_dict.get("ymax", 0)
            ]

            if "number" in lbl:
                fields["invoice_number"] = tag
            elif "date" in lbl:
                fields["invoice_date"] = tag
            elif "total" in lbl and "amount" in lbl:
                fields["total_amount"] = tag
            elif "order" in lbl:
                fields["buyer_name"] = tag

            # Add key words
            if key and key.get("tag"):
                k_bbox_dict = key.get("bbox", {})
                k_box = [
                    k_bbox_dict.get("xmin", 0),
                    k_bbox_dict.get("ymin", 0),
                    k_bbox_dict.get("xmax", 0),
                    k_bbox_dict.get("ymax", 0)
                ]
                for w in key["tag"].split():
                    ocr_words.append(OCRWord(text=w, bbox=k_box, conf=1.0))

            # Add value words
            if tag:
                for w in tag.split():
                    ocr_words.append(OCRWord(text=w, bbox=val_bbox, conf=1.0))

        rec = DocumentRecord(
            doc_id=doc_id,
            source_dataset="INV-CDIP",
            doc_class="invoice",
            image_path=dest_img_path,
            ocr_words=ocr_words,
            fields=fields
        )
        records.append(rec)

    return records