import io, json, os
import numpy as np
from PIL import Image
import gradio as gr
import cv2
import pandas as pd

from biz.segmentation import segment_cards
from biz.gemini import extract_from_crop
from biz.utils import crop_to_png_bytes, overlay_boxes, to_excel_file

# --- helpers ---
def np_to_pil(img_np) -> Image.Image:
    if isinstance(img_np, Image.Image): return img_np.convert("RGB")
    return Image.fromarray(img_np).convert("RGB")

def pil_to_bgr(pil: Image.Image):
    return cv2.cvtColor(np.array(pil), cv2.COLOR_RGB2BGR)

# --- Gradio functions ---
def do_segment(image_np):
    if image_np is None:
        return None, "[]", gr.update(visible=False), None
    pil = np_to_pil(image_np)
    bgr = pil_to_bgr(pil)
    boxes, W, H = segment_cards(bgr)
    overlay = overlay_boxes(pil, boxes)
    return overlay, json.dumps(boxes, ensure_ascii=False), gr.update(visible=True), None

def do_extract(image_np, boxes_json):
    if image_np is None or not boxes_json:
        return pd.DataFrame(), None
    pil = np_to_pil(image_np)
    try:
        boxes = json.loads(boxes_json)
    except Exception:
        boxes = []
    cards = []
    for b in boxes:
        crop = crop_to_png_bytes(pil, b["x"], b["y"], b["w"], b["h"])
        fields = extract_from_crop(crop, source_name="upload")
        fields["box_id"] = b["id"]
        cards.append(fields)
    df = pd.DataFrame(cards)
    xlsx_path = to_excel_file(cards)
    return df, xlsx_path

def clear_all():
    return None, None, "[]", gr.update(visible=False), pd.DataFrame(), None

# --- UI ---
with gr.Blocks(title="BizCards Extractor (Gradio)") as demo:
    gr.Markdown("## 💼 BizCards Extractor\nUpload → **Segment** → **Extract** → **Download Excel**")

    with gr.Row():
        with gr.Column(scale=3):
            in_img = gr.Image(type="numpy", label="Upload single or multi-card photo")
            with gr.Row():
                btn_seg = gr.Button("Segment", variant="primary")
                btn_ext = gr.Button("Extract", variant="secondary")
                btn_clear = gr.Button("Clear")
        with gr.Column(scale=2):
            out_img = gr.Image(label="Segmented preview (boxes)", interactive=False)
            out_table = gr.Dataframe(
                headers=["box_id","company","person_romaji","person_kanji","person_kana",
                         "title","department","email","phone","website","address_jp","notes","source_name"],
                wrap=True, height=350
            )
            dl = gr.File(label="Download Excel", visible=False)

    # hidden state for boxes in JSON
    boxes_state = gr.Textbox(label="boxes_json (debug)", visible=False, value="[]")

    # wiring
    btn_seg.click(fn=do_segment, inputs=[in_img],
                  outputs=[out_img, boxes_state, dl, dl])
    btn_ext.click(fn=do_extract, inputs=[in_img, boxes_state],
                  outputs=[out_table, dl])
    btn_clear.click(fn=clear_all, inputs=[],
                    outputs=[in_img, out_img, boxes_state, dl, out_table, dl])

    # show a warning if key missing
    if not os.getenv("GOOGLE_API_KEY") and not os.getenv("GOOGLE_GENAI_USE_VERTEXAI"):
        gr.Warning("GOOGLE_API_KEY is not set. Add it in Space → Settings → Variables & secrets.")

if __name__ == "__main__":
    demo.queue(max_size=16).launch()