# Imports
import os
from PIL import Image


class ImageTextDataset:
    def __init__(self, root):
        self.samples = []

        for folder in sorted(os.listdir(root)):
            sub = os.path.join(root, folder)
            img_dir = os.path.join(sub, "img")
            txt_dir = os.path.join(sub, "txt")

            if not (os.path.isdir(img_dir) and os.path.isdir(txt_dir)):
                continue

            for fname in sorted(os.listdir(img_dir)):
                if not fname.lower().endswith((".png", ".jpg", ".jpeg")):
                    continue

                img_path = os.path.join(img_dir, fname)

                # text file must have same name but .txt
                txt_name = os.path.splitext(fname)[0] + ".txt"
                txt_path = os.path.join(txt_dir, txt_name)

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

                self.samples.append((img_path, txt_path))
    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        img_path, txt_path = self.samples[idx]

        # Load image
        image = Image.open(img_path).convert("RGB")

        # Load text
        with open(txt_path, "r", encoding="utf-8") as f:
            text = f.read()

        return image, text

    def get_output_path(self, idx, output_subdir, debugset_root="/app/debugset"):
        """Get output path for saving OCR result to debugset folder.

        Args:
            idx: Sample index
            output_subdir: Subdirectory name (e.g., 'paddle_text', 'doctr_text')
            debugset_root: Root folder for debug output (default: /app/debugset)

        Returns:
            Path like /app/debugset/doc1/{output_subdir}/page_001.txt
        """
        img_path, _ = self.samples[idx]
        # img_path: /app/dataset/doc1/img/page_001.png
        # Extract relative path: doc1/img/page_001.png
        parts = img_path.split("/dataset/", 1)
        if len(parts) == 2:
            rel_path = parts[1]  # doc1/img/page_001.png
        else:
            rel_path = os.path.basename(img_path)

        # Replace /img/ with /{output_subdir}/
        rel_parts = rel_path.rsplit("/img/", 1)
        doc_folder = rel_parts[0]  # doc1
        fname = os.path.splitext(rel_parts[1])[0] + ".txt"  # page_001.txt

        out_dir = os.path.join(debugset_root, doc_folder, output_subdir)
        os.makedirs(out_dir, exist_ok=True)
        return os.path.join(out_dir, fname)