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

def convert_funsd(funsd_base_dir: str, target_dir: str, max_samples: int = None) -> List[DocumentRecord]:
    """
    Converts FUNSD form dataset into DocumentRecord format.
    """
    records = []
    os.makedirs(target_dir, exist_ok=True)

    # Search in training_data and testing_data
    subdirs = ["training_data", "testing_data"]
    total_loaded = 0

    for sdir in subdirs:
        anno_dir = os.path.join(funsd_base_dir, sdir, "annotations")
        img_dir = os.path.join(funsd_base_dir, sdir, "images")
        if not os.path.exists(anno_dir) or not os.path.exists(img_dir):
            continue

        for fname in sorted(os.listdir(anno_dir)):
            if not fname.endswith(".json"):
                continue
            if max_samples and total_loaded >= max_samples:
                break

            doc_id = f"funsd_{os.path.splitext(fname)[0]}"
            anno_path = os.path.join(anno_dir, fname)
            img_src = os.path.join(img_dir, os.path.splitext(fname)[0] + ".png")

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

            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:
                fdata = json.load(jf)

            fields = {k: None for k in CLASS_FIELD_SCHEMAS["application_form"]}
            ocr_words: List[OCRWord] = []
            qa_pairs = []

            # Extract words and detect key-value associations
            items = fdata.get("form", [])
            # Map item id to item
            id_to_item = {item.get("id"): item for item in items}

            for item in items:
                label = item.get("label", "")
                text = item.get("text", "").strip()
                words = item.get("words", [])

                for w in words:
                    w_text = w.get("text", "")
                    w_box = w.get("box", [0, 0, 0, 0])
                    if w_text:
                        ocr_words.append(OCRWord(text=w_text, bbox=w_box, conf=1.0))

                # Check question-answer linking
                if label == "question":
                    linked_ids = []
                    for link in item.get("linking", []):
                        if len(link) == 2 and link[0] == item.get("id"):
                            linked_ids.append(link[1])

                    answer_texts = [id_to_item[aid].get("text", "") for aid in linked_ids if aid in id_to_item]
                    ans_str = " ".join(answer_texts).strip() if answer_texts else None

                    q_lower = text.lower()
                    if ans_str:
                        qa_pairs.append(f"{text}: {ans_str}")
                        if any(k in q_lower for k in ["name", "applicant", "employee", "signature", "person"]) and not fields["applicant_name"]:
                            fields["applicant_name"] = ans_str
                        elif "date" in q_lower and not fields["date"]:
                            fields["date"] = ans_str
                        elif any(k in q_lower for k in ["address", "street", "city", "state", "location"]) and not fields["address"]:
                            fields["address"] = ans_str

            fields["other_key_value_pairs"] = len(qa_pairs)

            rec = DocumentRecord(
                doc_id=doc_id,
                source_dataset="FUNSD",
                doc_class="application_form",
                image_path=dest_img_path,
                ocr_words=ocr_words,
                fields=fields
            )
            records.append(rec)
            total_loaded += 1

        if max_samples and total_loaded >= max_samples:
            break

    return records