"""
Pizza Size Labeler — Gradio web app (Hugging Face Spaces)
===========================================================

Automatically loads photos from size.zip (bundled in this repo, extracted
once and cached). Click the matching size button for each photo, then
download the sorted result from the Output section.

All diagnostic messages go to the backend console / Space container logs
only (via the `logging` module) — nothing clutters the UI.
"""

import csv
import logging
import random
import shutil
import tempfile
import zipfile
from collections import Counter
from datetime import datetime
from pathlib import Path

import gradio as gr
from PIL import Image

try:
    import spaces
    _HAS_SPACES = True
except ImportError:
    _HAS_SPACES = False

if _HAS_SPACES:
    @spaces.GPU
    def _zerogpu_warmup():
        """No-op. This app never uses a GPU -- this function only exists so
        Hugging Face's ZeroGPU hardware tier finds a @spaces.GPU function at
        startup (it requires one, even if unused)."""
        return True

logging.basicConfig(level=logging.INFO, format='%(asctime)s [pizza_labeler] %(levelname)s: %(message)s')
log = logging.getLogger('pizza_labeler')

# ============================================================
# CONFIG
# ============================================================
SIZE_CLASSES = {
    '20': 'pizza_20sm',
    '25': 'pizza_25sm',
    '30': 'pizza_30sm',
    '35': 'pizza_35sm',
}
SKIP_LABEL = '__skipped__'
IMAGE_EXTS = {'.jpg', '.jpeg', '.png', '.bmp', '.webp'}
IGNORE_NAME_FRAGMENTS = {'__MACOSX', '.DS_Store', 'Thumbs.db'}

BUNDLED_ZIP_PATH = Path(__file__).resolve().parent / "size.zip"
SHARED_CACHE_DIR = Path(tempfile.gettempdir()) / "pizza_shared_cache"
_shared_bundle_cache = None  # (extract_dir: Path, images: list[Path]) once computed

# Output + manifest live outside any per-session temp dir so labeling
# progress survives page reloads / new browser sessions for as long as the
# container keeps running.
PERSISTENT_OUTPUT_ROOT = SHARED_CACHE_DIR / "output" / "size_labeled"
PERSISTENT_MANIFEST_PATH = PERSISTENT_OUTPUT_ROOT / "_manifest.csv"


# ============================================================
# Helpers
# ============================================================
def _is_junk(rel_path: Path) -> bool:
    return any(part in IGNORE_NAME_FRAGMENTS or part.startswith('.') for part in rel_path.parts)


def find_images_recursive(root: Path):
    found = []
    for p in root.rglob('*'):
        if p.is_file() and p.suffix.lower() in IMAGE_EXTS and not _is_junk(p.relative_to(root)):
            found.append(p)
    return sorted(found)


def dest_filename_for(path: Path, scan_root: Path) -> str:
    rel = path.relative_to(scan_root)
    if rel.parent == Path('.'):
        return path.name
    parent_tag = '_'.join(rel.parent.parts)
    return f"{parent_tag}__{path.name}"


def get_shared_bundled_images():
    """Extract BUNDLED_ZIP_PATH exactly once per running container and cache
    the result in memory for every subsequent session/page load."""
    global _shared_bundle_cache
    if _shared_bundle_cache is not None:
        return _shared_bundle_cache

    if not BUNDLED_ZIP_PATH.exists():
        raise FileNotFoundError(f"{BUNDLED_ZIP_PATH.name} not found next to app.py")

    extract_dir = SHARED_CACHE_DIR / "input"
    marker = SHARED_CACHE_DIR / ".extracted_ok"
    if not marker.exists():
        log.info(f"Extracting {BUNDLED_ZIP_PATH.name} ({BUNDLED_ZIP_PATH.stat().st_size / 1e6:.1f} MB) ...")
        extract_dir.mkdir(parents=True, exist_ok=True)
        with zipfile.ZipFile(BUNDLED_ZIP_PATH, 'r') as zf:
            zf.extractall(extract_dir)
        marker.touch()
        log.info("Extraction complete.")
    else:
        log.info("Using previously extracted cache (no re-extraction needed).")

    images = find_images_recursive(extract_dir)
    log.info(f"Found {len(images)} image file(s).")
    _shared_bundle_cache = (extract_dir, images)
    return _shared_bundle_cache


def append_manifest(manifest_path: Path, key: str, label: str):
    is_new = not manifest_path.exists()
    with open(manifest_path, 'a', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        if is_new:
            writer.writerow(['filename', 'label', 'timestamp'])
        writer.writerow([key, label, datetime.now().isoformat(timespec='seconds')])


def remove_last_manifest_entry(manifest_path: Path, key: str):
    if not manifest_path.exists():
        return
    with open(manifest_path, newline='', encoding='utf-8') as f:
        rows = list(csv.reader(f))
    if not rows:
        return
    header, body = rows[0], rows[1:]
    for i in range(len(body) - 1, -1, -1):
        if body[i][0] == key:
            del body[i]
            break
    with open(manifest_path, 'w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        writer.writerow(header)
        writer.writerows(body)


def load_manifest(manifest_path: Path):
    counts = {}
    if manifest_path.exists():
        with open(manifest_path, newline='', encoding='utf-8') as f:
            for row in csv.DictReader(f):
                counts[row['filename']] = row['label']
    return counts


def summary_text(manifest_path: Path):
    counts = Counter(load_manifest(manifest_path).values())
    parts = [f"{key}: {counts.get(key, 0)}" for key in SIZE_CLASSES]
    parts.append(f"skipped: {counts.get(SKIP_LABEL, 0)}")
    return "Totals — " + "   ".join(parts)


def fresh_state():
    return {
        'input_root': None,
        'output_root': None,
        'manifest_path': None,
        'images': [],   # list of (abs_path_str, rel_key_str)
        'idx': 0,
        'history': [],  # list of (abs_path_str, rel_key_str, label, dest_path_str_or_None)
        'total_images': 0,
    }


# ============================================================
# Core actions
# ============================================================
def load_photos(state):
    state = fresh_state()

    output_root = PERSISTENT_OUTPUT_ROOT
    output_root.mkdir(parents=True, exist_ok=True)
    for folder in SIZE_CLASSES.values():
        (output_root / folder).mkdir(parents=True, exist_ok=True)

    state['output_root'] = str(output_root)
    state['manifest_path'] = str(PERSISTENT_MANIFEST_PATH)

    try:
        input_root, all_images = get_shared_bundled_images()
    except Exception as e:
        log.error(f"Could not load bundled zip: {e}")
        return None, "Could not load photos — see server logs.", "", state

    state['input_root'] = str(input_root)

    if not all_images:
        log.warning("No images found in size.zip.")
        return None, "No images found.", "", state

    labeled_map = load_manifest(Path(state['manifest_path']))
    pairs = [(str(p), str(p.relative_to(input_root).as_posix())) for p in all_images]
    remaining = [(p, k) for p, k in pairs if k not in labeled_map]

    random.Random(42).shuffle(remaining)
    state['images'] = remaining
    state['idx'] = 0
    state['history'] = []
    state['total_images'] = len(pairs)

    log.info(f"Session ready with {len(remaining)} photo(s) to label.")

    if not remaining:
        return None, "All photos already labeled.", summary_text(Path(state['manifest_path'])), state

    first_path, first_key = remaining[0]
    progress = f"labeled {len(labeled_map)} / {len(pairs)}   —   1 / {len(remaining)} remaining   ({first_key})"
    return first_path, progress, summary_text(Path(state['manifest_path'])), state


def _current_view(state):
    manifest_path = Path(state['manifest_path'])
    idx = state['idx']
    images = state['images']
    total = state.get('total_images', len(images))
    labeled_so_far = total - len(images) + idx
    if idx >= len(images):
        return None, f"All done! {labeled_so_far} / {total} labeled.", summary_text(manifest_path), state
    abs_path, key = images[idx]
    progress = f"labeled {labeled_so_far} / {total}   —   {idx + 1} / {len(images)} remaining   ({key})"
    return abs_path, progress, summary_text(manifest_path), state


def _advance(state, label_key):
    manifest_path = Path(state['manifest_path'])
    input_root = Path(state['input_root'])
    output_root = Path(state['output_root'])

    idx = state['idx']
    images = state['images']
    if idx >= len(images):
        return _current_view(state)

    abs_path, key = images[idx]
    path = Path(abs_path)

    if label_key == SKIP_LABEL:
        dest_path = None
        log.info(f"Skipped: {key}")
    else:
        folder = SIZE_CLASSES[label_key]
        dest_name = dest_filename_for(path, input_root)
        dest_path = output_root / folder / dest_name
        shutil.copy2(path, dest_path)
        log.info(f'Labeled "{key}" -> {label_key} ({folder}/{dest_name})')

    append_manifest(manifest_path, key, label_key)
    state['history'].append((abs_path, key, label_key, str(dest_path) if dest_path else None))
    state['idx'] += 1

    return _current_view(state)


def label_20(state):
    return _advance(state, '20')


def label_25(state):
    return _advance(state, '25')


def label_30(state):
    return _advance(state, '30')


def label_35(state):
    return _advance(state, '35')


def skip(state):
    return _advance(state, SKIP_LABEL)


def undo(state):
    if not state.get('history'):
        log.warning("Nothing to undo.")
        return _current_view(state)

    abs_path, key, label_key, dest_path = state['history'].pop()
    if dest_path:
        p = Path(dest_path)
        if p.exists():
            p.unlink()
    remove_last_manifest_entry(Path(state['manifest_path']), key)
    state['idx'] -= 1
    log.info(f"Undid label for: {key}")
    return _current_view(state)


def prepare_download(state):
    if not state.get('output_root'):
        log.warning("Nothing to download yet.")
        return None
    output_root = Path(state['output_root'])
    work_dir = Path(tempfile.mkdtemp(prefix='pizza_zip_'))
    zip_path = work_dir / 'pizza_labeled_output.zip'
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
        for f in output_root.rglob('*'):
            if f.is_file():
                zf.write(f, f.relative_to(output_root.parent))
    log.info(f"Prepared download: {zip_path.name}")
    return str(zip_path)


# ============================================================
# UI
# ============================================================
with gr.Blocks(title="Pizza Size Labeler") as demo:
    state = gr.State(fresh_state())

    gr.Markdown(
        "## 🍕 Pizza Size Labeler\n"
        "Click the size that matches the photo (20 / 25 / 30 / 35)."
    )

    progress_box = gr.Textbox(label="Progress", interactive=False)
    image_view = gr.Image(label="Current photo", type="filepath", height=420)

    with gr.Row():
        b20 = gr.Button("20")
        b25 = gr.Button("25")
        b30 = gr.Button("30")
        b35 = gr.Button("35")
    with gr.Row():
        skip_btn = gr.Button("Skip")
        undo_btn = gr.Button("Undo")

    summary_box = gr.Textbox(label="Totals", interactive=False)

    with gr.Row():
        download_btn = gr.Button("Prepare download zip ⬇", variant="primary")
        download_file = gr.File(label="Labeled output (zip)", interactive=False)

    demo.load(load_photos, inputs=[state], outputs=[image_view, progress_box, summary_box, state])

    for btn, fn in [(b20, label_20), (b25, label_25), (b30, label_30), (b35, label_35),
                    (skip_btn, skip), (undo_btn, undo)]:
        btn.click(fn, inputs=[state], outputs=[image_view, progress_box, summary_box, state])

    download_btn.click(prepare_download, inputs=[state], outputs=[download_file])


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