import gradio as gr
from huggingface_hub import HfApi, hf_hub_download, upload_file
import os
from PIL import Image
import random

# === CONFIGURATION ===
HF_TOKEN = os.environ.get("HF_TOKEN")
HF_IMAGE_REPO = "saleh-c4/arabic-ocr-images"
HF_LABEL_REPO = "saleh-c4/arabic-ocr-labels"
TEMP_DIR = "temp_downloads"

api = HfApi()

# === TRACKING STATE ===
all_images = []  # all line images with page
current_image = None

# === INITIALIZE IMAGE LIST ===
def fetch_all_images():
    global all_images
    files = api.list_repo_files(HF_IMAGE_REPO, repo_type="dataset", token=HF_TOKEN)
    all_images = sorted([
        f for f in files if f.endswith((".png", ".jpg", ".jpeg")) and "/" in f
    ])

fetch_all_images()

# === CHECK IF IMAGE ALREADY LABELED ===
def is_labeled(image_path):
    filename = os.path.basename(image_path)
    page = image_path.split("/")[0]
    label_file = f"{page}.txt"

    try:
        label_path = hf_hub_download(
            repo_id=HF_LABEL_REPO,
            repo_type="dataset",
            filename=label_file,
            token=HF_TOKEN,
            local_dir=TEMP_DIR,
            local_dir_use_symlinks=False
        )
        with open(label_path, "r", encoding="utf-8") as f:
            return any(filename in line for line in f)
    except:
        return False

# === LOAD A RANDOM UNLABELED IMAGE ===
def load_image():
    global current_image

    attempts = 0
    while attempts < 20:
        candidate = random.choice(all_images)
        if not is_labeled(candidate):
            current_image = candidate
            break
        attempts += 1
    else:
        return None, "", "✅ كل الصور مُعنونة!"

    local_path = hf_hub_download(
        repo_id=HF_IMAGE_REPO,
        repo_type="dataset",
        filename=current_image,
        token=HF_TOKEN,
        local_dir=TEMP_DIR,
        local_dir_use_symlinks=False
    )
    img = Image.open(local_path)
    status = f"{current_image} (لم تُعنون بعد)"
    return img, "", status

# === HANDLE LABEL ===
def handle_label(text):
    global current_image

    if not current_image:
        return load_image()

    filename = os.path.basename(current_image)
    page = current_image.split("/")[0]
    label_file = f"{page}.txt"
    local_label_path = f"{TEMP_DIR}/{label_file}"

    try:
        # Try downloading the existing label file
        existing_lines = []
        try:
            label_path = hf_hub_download(
                repo_id=HF_LABEL_REPO,
                repo_type="dataset",
                filename=label_file,
                token=HF_TOKEN,
                local_dir=TEMP_DIR,
                local_dir_use_symlinks=False
            )
            with open(label_path, "r", encoding="utf-8") as f:
                existing_lines = [line.strip() for line in f if line.strip()]
        except:
            pass  # no existing file

        # Add the new line
        new_line = f"{filename}\t{text.strip()}"
        existing_lines.append(new_line)

        # Sort lines by filename
        sorted_lines = sorted(existing_lines, key=lambda x: x.split("\t")[0])

        # Save back
        with open(local_label_path, "w", encoding="utf-8") as f:
            f.write("\n".join(sorted_lines) + "\n")

        upload_file(
            path_or_fileobj=local_label_path,
            path_in_repo=label_file,
            repo_id=HF_LABEL_REPO,
            repo_type="dataset",
            token=HF_TOKEN
        )

        if os.path.exists(local_label_path):
            os.remove(local_label_path)

        print(f"✅ Labeled and sorted: {filename} → {label_file}")
    except Exception as e:
        print(f"❌ Failed to label: {e}")

    return load_image()


# === UI ===
def setup():
    img, _, desc = load_image()
    return gr.update(value=img), gr.update(value="", placeholder="اكتب التسمية هنا..."), gr.update(value=desc)

with gr.Blocks(css="""
#label_image label {
    display: none !important;
}
#label_image .absolute.top-0.right-0 {
    display: none !important;  /* Optional: hide buttons */
}
#label_image canvas,
#label_image img {
    max-width: 90% !important;
    margin-top: 20px;
    margin-bottom: 10px;  /* creates space from bottom buttons */
    border-radius: 10px;
    box-shadow: 0 0 5px rgba(0,0,0,0.2);
}
""") as demo:

    gr.Markdown("🚨 <b>ملاحظة:</b> تسمية الصورة بـ <code>0</code> تعني أنها لا تحتوي على نص وسيتم حذفها من قاعدة البيانات.")
    img = gr.Image(elem_id="label_image", container=False)
    textbox = gr.Textbox(label="التسمية", rtl=True)
    status = gr.Textbox(label="الوصف", interactive=False)

    textbox.submit(fn=handle_label, inputs=textbox, outputs=[img, textbox, status])
    demo.load(fn=setup, outputs=[img, textbox, status])

if __name__ == "__main__":
    demo.launch()