import os
import cv2
import json
import base64
import uuid
import tempfile
import numpy as np
import gradio as gr
from concurrent.futures import ThreadPoolExecutor, as_completed
from roboflow import Roboflow
from openai import OpenAI

# ============================================================
# CONFIGURATION & COLOR REFERENCES
# ============================================================
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
ROBOFLOW_API_KEY = "o8GqfxLrU6X4RzkoageL"

ROBOFLOW_PROJECT = "terminal-segmentation-uqf2e"
ROBOFLOW_VERSION = 3
CONFIDENCE_THRESHOLD = 40

GPT_MODEL = "gpt-4o"          # switch to "gpt-4o-mini" for extra speed if accuracy still holds
MAX_WORKERS = 8               # parallel GPT calls -> tune to your OpenAI account's rate limit
CROP_PADDING = 6              # px of padding added around every detected wire box before cropping
OCR_CROP_PADDING = 10         # slightly larger padding for number tags so characters are never clipped
OCR_TARGET_HEIGHT = 180       # tiny number-tag crops are upscaled to this height before OCR
GPT_SEED = 42                 # best-effort reproducibility for GPT-4o across repeated runs
                               # (OpenAI does not guarantee perfect determinism even at temperature=0,
                               # but a fixed seed noticeably reduces run-to-run output variance)

# Initialize clients once (module load time)
rf = Roboflow(api_key=ROBOFLOW_API_KEY)
project = rf.workspace().project(ROBOFLOW_PROJECT)
model = project.version(ROBOFLOW_VERSION).model
openai_client = OpenAI(api_key=OPENAI_API_KEY)

# --- ADJUSTED SATURATION THRESHOLD ---
# Lowered from 50 to 25 so that shaded or dusty colored wires are not misclassified as grey.
ACHROMATIC_SAT_THRESHOLD = 25   
BLACK_VALUE_MAX = 50
WHITE_VALUE_MIN = 195
BROWN_VALUE_MAX = 120          # orange-hue pixel darker than this -> "brown" instead of "orange"

HUE_BANDS = [
    (8,   "red"),
    (20,  "orange_or_brown"),   # resolved to orange/brown based on value, see map_hsv_to_name
    (35,  "yellow"),
    (85,  "green"),
    (135, "blue"),
    (160, "violet"),
    (180, "pink"),
]

# --- OPTIMIZED COLOR REFERENCE DICTIONARY ---
# Adjusted Hue, Saturation, and Value targets to match the true physical wires
COLOR_REFERENCE_HSV = {
    "yellow": (25, 200, 180),    # Vibrant terminal yellow
    "green": (64, 210, 160),     # Earth green
    "blue": (115, 230, 180),
    "red": (0, 220, 200),
    "orange": (14, 220, 220),
    "brown": (8, 150, 60),       # Kept brown's hue low to prevent overlap with yellow
    "black": (0, 0, 25),
    "white": (0, 0, 240),
    "grey": (0, 0, 120),
}

# --- OCR CHARACTER-CONFUSION GROUPS ---
# Characters that look alike on small/blurry crops. Used both to warn GPT-4o in the
# prompt, and for a post-OCR consensus correction pass across the whole sheet.
CONFUSABLE_GROUPS = [
    {"0", "O"},
    {"1", "I", "L"},
    {"5", "S"},
    {"8", "B"},
    {"2", "Z"},
    {"4", "A"},   # e.g. "43-M2" misread as "A3-M2"
    {"Y", "V"},   # e.g. "Y180" misread as "V180"
]
CONFUSABLE_MAP = {}
for _group in CONFUSABLE_GROUPS:
    for _ch in _group:
        CONFUSABLE_MAP[_ch] = _group


# ============================================================
# IMAGE PREPROCESSING
# ============================================================
def preprocess_image(image_bgr):
    """Sharpen and balance lighting to optimize OCR and color recognition."""
    img_yuv = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2YUV)
    img_yuv[:, :, 0] = cv2.equalizeHist(img_yuv[:, :, 0])
    enhanced = cv2.cvtColor(img_yuv, cv2.COLOR_YUV2BGR)
    kernel = np.array([[0, -0.5, 0], [-0.5, 3, -0.5], [0, -0.5, 0]])
    return cv2.filter2D(enhanced, -1, kernel)


def is_bbox_inside_polygon(bbox, polygon_points):
    """Determines if a bounding box center falls within the segment's polygon."""
    if not polygon_points or len(polygon_points) < 3:
        return False
    poly_array = np.array([[p['x'], p['y']] for p in polygon_points], dtype=np.int32)
    cx, cy = int(bbox['x']), int(bbox['y'])
    return cv2.pointPolygonTest(poly_array, (cx, cy), False) >= 0


def assign_items_to_tracks(items, tracks):
    """
    Assign each detected wire/number to EXACTLY ONE terminal column (never more than one).

    Terminal columns sit tightly packed side by side, so their horizontal x-ranges can
    overlap slightly, and wires often bend sideways on their way to a tag. Testing "is this
    item inside column N's range" independently per column (the previous approach) let a
    single item pass that test for two neighboring columns at once -- the item would get
    duplicated into one column while silently disappearing (reported as Blank) from its
    true column, and which column "won" could vary between runs.

    This function fixes that by making a single best-column decision per item:
      1. If the item falls inside one or more terminal polygons, pick the polygon whose
         column center (x) is closest to the item (handles the wire's actual bent path).
      2. Otherwise, fall back to the terminal column whose center (x) is nearest overall
         (handles tags/wires that legitimately sit outside a tight segmentation polygon).
    Every item ends up in exactly one column's list, eliminating cross-column bleed.
    """
    track_x = [t['x'] for t in tracks]
    assignment = [[] for _ in tracks]

    for item in items:
        containing = [
            i for i, t in enumerate(tracks)
            if t.get('points') and is_bbox_inside_polygon(item, t['points'])
        ]
        if containing:
            best_idx = min(containing, key=lambda i: abs(track_x[i] - item['x']))
        else:
            best_idx = min(range(len(tracks)), key=lambda i: abs(track_x[i] - item['x']))
        assignment[best_idx].append(item)

    return assignment


def get_bbox_coords(item):
    if not item:
        return None
    return {
        "x_min": int(item['x'] - item['width'] / 2),
        "y_min": int(item['y'] - item['height'] / 2),
        "x_max": int(item['x'] + item['width'] / 2),
        "y_max": int(item['y'] + item['height'] / 2),
    }


# ============================================================
# CROPPING HELPERS
# ============================================================
def crop_region(image, box, pad=CROP_PADDING):
    if box is None:
        return None
    h, w = image.shape[:2]
    x1 = max(0, box["x_min"] - pad)
    y1 = max(0, box["y_min"] - pad)
    x2 = min(w, box["x_max"] + pad)
    y2 = min(h, box["y_max"] + pad)
    if x2 <= x1 or y2 <= y1:
        return None
    return image[y1:y2, x1:x2].copy()


def prepare_crop_for_ocr(crop):
    """Upscale small number-tag crops (standard variant) so GPT-4o can read the text reliably."""
    if crop is None or crop.size == 0:
        return None
    h, w = crop.shape[:2]
    if h == 0 or w == 0:
        return None
    scale = min(OCR_TARGET_HEIGHT / float(h), 5.0)
    new_w, new_h = max(1, int(w * scale)), max(1, int(h * scale))
    return cv2.resize(crop, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)


def enhance_crop_for_ocr(crop):
    """
    Produces a second, contrast-enhanced + sharpened variant of a number-tag crop.
    GPT-4o is given BOTH the standard and enhanced variant of the same tag so it can
    cross-check ambiguous characters (e.g. Y vs V, 4 vs A, 0 vs O, 1 vs I, S vs 5, B vs 8)
    against two different renderings instead of guessing from a single blurry read.
    """
    if crop is None or crop.size == 0:
        return None
    h, w = crop.shape[:2]
    if h == 0 or w == 0:
        return None

    # CLAHE (local contrast enhancement) on the L channel to make printed characters
    # stand out clearly from the white sleeve background, even under uneven lighting.
    lab = cv2.cvtColor(crop, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    l = clahe.apply(l)
    enhanced = cv2.cvtColor(cv2.merge((l, a, b)), cv2.COLOR_LAB2BGR)

    # Stronger unsharp mask specifically tuned for thin printed text edges
    blurred = cv2.GaussianBlur(enhanced, (0, 0), sigmaX=1.2)
    sharpened = cv2.addWeighted(enhanced, 1.6, blurred, -0.6, 0)

    scale = min(OCR_TARGET_HEIGHT / float(h), 5.0)
    new_w, new_h = max(1, int(w * scale)), max(1, int(h * scale))
    return cv2.resize(sharpened, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)


def encode_b64(img_bgr):
    if img_bgr is None:
        return None
    ok, buf = cv2.imencode(".png", img_bgr)
    if not ok:
        return None
    return base64.b64encode(buf).decode("utf-8")


# ============================================================
# LOCAL, DETERMINISTIC WIRE-COLOR CLASSIFICATION
# ============================================================
def classify_wire_color(crop_bgr):
    if crop_bgr is None or crop_bgr.size == 0:
        return "unknown"
    hsv = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2HSV)
    pixels = hsv.reshape(-1, 3).astype(np.float32)

    # Exclude reflections (highly bright/fully white specular highlights) and very deep dark shadows
    s = pixels[:, 1]
    v = pixels[:, 2]
    mask = (v > 30) & (v < 240)
    filtered = pixels[mask] if mask.sum() > 20 else pixels

    # Calculate median metrics safely
    median_hsv = np.median(filtered, axis=0)
    median_s = median_hsv[1]
    median_v = median_hsv[2]

    # 1. Deterministic check for neutral/achromatic tones (removes white/grey false positives on colored wires)
    if median_s < ACHROMATIC_SAT_THRESHOLD:
        if median_v < BLACK_VALUE_MAX:
            return "black"
        elif median_v > WHITE_VALUE_MIN:
            return "white"
        else:
            return "grey"

    # 2. ACCURATE YELLOW OVERRIDE RULE
    # OpenCV Hue for pure Yellow runs between 18 and 42. By explicitly checking this range
    # first, we ensure shaded yellow wires are never mismatched to adjacent brown references.
    if 18 <= median_hsv[0] <= 42:
        return "yellow"

    # 3. Fallback to Weighted distance matching for all other colors
    best_name, best_dist = "unknown", float("inf")
    for name, ref in COLOR_REFERENCE_HSV.items():
        # Only evaluate chromatic profiles for chromatic detections
        if name in ["white", "grey", "black", "yellow"]:
            continue
        # Hue distance on 180-deg circle
        dh = min(abs(median_hsv[0] - ref[0]), 180 - abs(median_hsv[0] - ref[0]))
        ds = abs(median_hsv[1] - ref[1])
        dv = abs(median_hsv[2] - ref[2])

        # Heavy weight to Hue, moderate to Saturation, low to Value
        dist = (dh * 3.0) ** 2 + (ds * 0.8) ** 2 + (dv * 0.2) ** 2
        if dist < best_dist:
            best_dist, best_name = dist, name

    return best_name


# ============================================================
# GPT-4o OCR — ONE small call PER COLUMN, run in parallel
# Each tag is sent as TWO variants (standard + contrast-enhanced) so GPT-4o can
# cross-check its reading instead of committing to a single ambiguous render.
# ============================================================
def ocr_column_numbers(column_index, top_std_b64, top_enh_b64, bottom_std_b64, bottom_enh_b64):
    if top_std_b64 is None and bottom_std_b64 is None:
        return {"column": column_index, "top_text": "", "bottom_text": ""}

    content = [
        {"type": "text", "text": (
            "You are reading small cropped photos of white wire-marker sleeve tags used on "
            "electrical terminal blocks. For each tag, you are given TWO images of the SAME "
            "tag: 'Version A' (standard render) and 'Version B' (contrast-enhanced render). "
            "Cross-check both versions letter by letter before deciding the final text.\n\n"
            "Be extremely careful with visually similar characters that are commonly confused "
            "in this font, especially on low-resolution crops:\n"
            "  - 'Y' vs 'V' (Y has a straight vertical stem below the join; V has no stem, "
            "it is a clean pointed checkmark shape all the way to the bottom)\n"
            "  - '4' vs 'A' (4 has a flat horizontal crossbar and an open top; A is a closed "
            "triangle/peak at the top with a crossbar lower down — do not read a printed '4' as 'A')\n"
            "  - '0' (zero) vs 'O' (letter O)\n"
            "  - '1' vs 'I' vs 'L'\n"
            "  - '8' vs 'B', '5' vs 'S', '2' vs 'Z'\n"
            "  - a hyphen '-' vs no character at all (do not insert a hyphen unless clearly printed)\n\n"
            "Preserve the exact characters printed, including hyphens (e.g. distinguish 'ED' vs "
            "'H-ED', 'Y180' vs 'V180', and '43-M2' vs 'A3-M2'). If a tag shows no legible printed "
            "text, or is blank/not present, return an empty string \"\" for that field — never "
            "guess a value you are not confident about.\n\n"
            "Respond ONLY with strict JSON: {\"top_text\": \"...\", \"bottom_text\": \"...\"}"
        )}
    ]

    if top_std_b64:
        content.append({"type": "text", "text": "TOP tag — Version A:"})
        content.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{top_std_b64}"}})
        if top_enh_b64:
            content.append({"type": "text", "text": "TOP tag — Version B (enhanced):"})
            content.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{top_enh_b64}"}})
    else:
        content.append({"type": "text", "text": "TOP tag: not detected -> top_text must be \"\""})

    if bottom_std_b64:
        content.append({"type": "text", "text": "BOTTOM tag — Version A:"})
        content.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{bottom_std_b64}"}})
        if bottom_enh_b64:
            content.append({"type": "text", "text": "BOTTOM tag — Version B (enhanced):"})
            content.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{bottom_enh_b64}"}})
    else:
        content.append({"type": "text", "text": "BOTTOM tag: not detected -> bottom_text must be \"\""})

    try:
        response = openai_client.chat.completions.create(
            model=GPT_MODEL,
            response_format={"type": "json_object"},
            messages=[{"role": "user", "content": content}],
            max_tokens=200,
            temperature=0,
            seed=GPT_SEED,
        )
        result = json.loads(response.choices[0].message.content)
        return {
            "column": column_index,
            "top_text": (result.get("top_text") or "").strip(),
            "bottom_text": (result.get("bottom_text") or "").strip(),
        }
    except Exception as e:
        return {"column": column_index, "top_text": "", "bottom_text": "", "error": str(e)}


# ============================================================
# CROSS-COLUMN CONSENSUS CORRECTION (post-OCR)
# Terminal sheets typically reuse the same numeric/letter prefixes across many
# columns (e.g. "43-M1", "43-A1", "43-M2", "43-A2"). If one isolated reading only
# differs from a much more common reading elsewhere on the SAME sheet by a single
# commonly-confused character (4/A, Y/V, 0/O, ...), it is very likely a misread and
# gets corrected to the dominant, more common variant. This is conservative: it only
# fires when the alternative is clearly more common (seen at least 2x more often),
# so it will not "invent" corrections on sheets with genuinely unique tags.
# ============================================================
def _generate_confusable_variants(text_upper):
    variants = set()
    chars = list(text_upper)
    for i, ch in enumerate(chars):
        alternates = CONFUSABLE_MAP.get(ch)
        if not alternates:
            continue
        for alt in alternates:
            if alt == ch:
                continue
            new_chars = chars.copy()
            new_chars[i] = alt
            variants.add("".join(new_chars))
    return variants


def apply_consensus_correction(ocr_results):
    # Build a frequency table of every non-blank tag text seen across the whole sheet
    freq = {}
    for res in ocr_results.values():
        for key in ("top_text", "bottom_text"):
            t = (res.get(key) or "").strip()
            if t:
                t_up = t.upper()
                freq[t_up] = freq.get(t_up, 0) + 1

    for res in ocr_results.values():
        for key in ("top_text", "bottom_text"):
            t = (res.get(key) or "").strip()
            if not t:
                continue
            t_up = t.upper()
            current_count = freq.get(t_up, 0)
            best_variant, best_count = t_up, current_count
            for variant in _generate_confusable_variants(t_up):
                vc = freq.get(variant, 0)
                if vc > best_count:
                    best_variant, best_count = variant, vc
            # Only correct when the alternative is clearly dominant on this sheet
            if best_variant != t_up and best_count >= current_count + 2:
                res[key] = best_variant

    return ocr_results


# ============================================================
# ANNOTATION
# ============================================================
def draw_annotations(image_bgr, tracks, wires, numbers):
    vis = image_bgr.copy()
    for t in tracks:
        pts = t.get('points')
        if pts:
            poly = np.array([[int(p['x']), int(p['y'])] for p in pts], dtype=np.int32)
            cv2.polylines(vis, [poly], isClosed=True, color=(0, 255, 255), thickness=2)
        else:
            x1, y1 = int(t['x'] - t['width'] / 2), int(t['y'] - t['height'] / 2)
            x2, y2 = int(t['x'] + t['width'] / 2), int(t['y'] + t['height'] / 2)
            cv2.rectangle(vis, (x1, y1), (x2, y2), (0, 255, 255), 2)
    for w in wires:
        x1, y1 = int(w['x'] - w['width'] / 2), int(w['y'] - w['height'] / 2)
        x2, y2 = int(w['x'] + w['width'] / 2), int(w['y'] + w['height'] / 2)
        cv2.rectangle(vis, (x1, y1), (x2, y2), (255, 100, 0), 2)
    for n in numbers:
        x1, y1 = int(n['x'] - n['width'] / 2), int(n['y'] - n['height'] / 2)
        x2, y2 = int(n['x'] + n['width'] / 2), int(n['y'] + n['height'] / 2)
        cv2.rectangle(vis, (x1, y1), (x2, y2), (0, 200, 0), 2)
    return cv2.cvtColor(vis, cv2.COLOR_BGR2RGB)


# ============================================================
# MAIN PIPELINE
# ============================================================
def process_and_verify(image_input):
    if image_input is None:
        return None, "### Error: Please upload an image first."
    if not OPENAI_API_KEY:
        return None, "### ❌ Error: `OPENAI_API_KEY` environment variable is missing."

    image_bgr = cv2.cvtColor(image_input, cv2.COLOR_RGB2BGR)
    processed_img = preprocess_image(image_bgr)

    # Unique temp filename -> safe for multiple concurrent Gradio users
    temp_input_path = os.path.join(tempfile.gettempdir(), f"terminal_{uuid.uuid4().hex}.jpg")
    cv2.imwrite(temp_input_path, processed_img)

    try:
        prediction_response = model.predict(temp_input_path, confidence=CONFIDENCE_THRESHOLD)
        predictions = prediction_response.json().get('predictions', [])
    finally:
        if os.path.exists(temp_input_path):
            os.remove(temp_input_path)

    tracks = sorted(
        [p for p in predictions if "terminal-segmentation" in p['class'].lower()],
        key=lambda k: k['x']
    )
    all_wires = [p for p in predictions if p['class'].lower() == "wire"]
    all_numbers = [p for p in predictions if "number" in p['class'].lower()]

    ui_display_image = draw_annotations(processed_img, tracks, all_wires, all_numbers)

    if not tracks:
        return ui_display_image, "### ❌ Error: No active terminal tracks detected by the model."

    # ---- Build per-column crop metadata ----
    # Each wire / number is assigned to exactly one column up front (see
    # assign_items_to_tracks) so no item can ever bleed into two neighboring columns.
    wires_by_track = assign_items_to_tracks(all_wires, tracks)
    numbers_by_track = assign_items_to_tracks(all_numbers, tracks)

    columns = []
    for index, track in enumerate(tracks):
        slot_wires = wires_by_track[index]
        slot_numbers = numbers_by_track[index]

        top_w = sorted([w for w in slot_wires if w['y'] < track['y']], key=lambda k: k['y'])
        bot_w = sorted([w for w in slot_wires if w['y'] >= track['y']], key=lambda k: k['y'], reverse=True)
        top_n = sorted([n for n in slot_numbers if n['y'] < track['y']], key=lambda k: k['y'])
        bot_n = sorted([n for n in slot_numbers if n['y'] >= track['y']], key=lambda k: k['y'], reverse=True)

        top_wire_box = get_bbox_coords(top_w[0] if top_w else None)
        bot_wire_box = get_bbox_coords(bot_w[0] if bot_w else None)
        top_num_box = get_bbox_coords(top_n[0] if top_n else None)
        bot_num_box = get_bbox_coords(bot_n[0] if bot_n else None)

        # Crop color regions directly from pristine ORIGINAL image_bgr
        top_wire_crop = crop_region(image_bgr, top_wire_box)
        bot_wire_crop = crop_region(image_bgr, bot_wire_box)

        # Text OCR crops use the sharpened processed_img, with extra padding so
        # characters near the edge of the detection box are never clipped.
        top_num_crop_raw = crop_region(processed_img, top_num_box, pad=OCR_CROP_PADDING)
        bot_num_crop_raw = crop_region(processed_img, bot_num_box, pad=OCR_CROP_PADDING)

        top_num_std = prepare_crop_for_ocr(top_num_crop_raw)
        top_num_enh = enhance_crop_for_ocr(top_num_crop_raw)
        bot_num_std = prepare_crop_for_ocr(bot_num_crop_raw)
        bot_num_enh = enhance_crop_for_ocr(bot_num_crop_raw)

        columns.append({
            "column": index + 1,
            "top_color": classify_wire_color(top_wire_crop),
            "bottom_color": classify_wire_color(bot_wire_crop),
            "top_num_std_b64": encode_b64(top_num_std),
            "top_num_enh_b64": encode_b64(top_num_enh),
            "bottom_num_std_b64": encode_b64(bot_num_std),
            "bottom_num_enh_b64": encode_b64(bot_num_enh),
        })

    # ---- Parallel GPT-4o OCR calls, one small call per column ----
    ocr_results = {}
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = {
            executor.submit(
                ocr_column_numbers,
                c["column"],
                c["top_num_std_b64"], c["top_num_enh_b64"],
                c["bottom_num_std_b64"], c["bottom_num_enh_b64"],
            ): c["column"]
            for c in columns
        }
        for future in as_completed(futures):
            res = future.result()
            ocr_results[res["column"]] = res

    # ---- Cross-column consensus correction (fixes isolated confusable misreads,
    #      e.g. a lone "A3-M2" when "43-M1" / "43-A1" / "43-A2" already confirm "43-") ----
    ocr_results = apply_consensus_correction(ocr_results)

    # ---- Build report ----
    report = "## 📊 Optimized GPT-4o Verification Report\n"
    report += f"- **Verified Columns:** {len(columns)}\n\n"
    report += "| Column | Top Tag | Bottom Tag | Top Color | Bottom Color | Status | Reason |\n"
    report += "| :---: | :---: | :---: | :---: | :---: | :---: | :--- |\n"

    for c in columns:
        col_num = c["column"]
        ocr = ocr_results.get(col_num, {"top_text": "", "bottom_text": ""})
        t_text = ocr.get("top_text", "").strip()
        b_text = ocr.get("bottom_text", "").strip()
        t_color = c["top_color"]
        b_color = c["bottom_color"]

        both_blank = (t_text == "") and (b_text == "")
        one_blank = (t_text == "") != (b_text == "")
        text_matches = (not both_blank) and (not one_blank) and (t_text.lower() == b_text.lower())
        color_matches = (t_color != "unknown") and (t_color == b_color)

        reasons = []
        if both_blank:
            status = "⚠️ NO TAG"
            reasons.append("No number tag detected on either wire (informational, not a mismatch)")
        elif one_blank:
            status = "❌ Batsu"
            reasons.append("Top tag blank" if t_text == "" else "Bottom tag blank")
        elif not text_matches:
            status = "❌ Batsu"
            reasons.append(f"Text mismatch ('{t_text}' vs '{b_text}')")
        elif not color_matches:
            status = "❌ Batsu"
            reasons.append(f"Color mismatch ('{t_color}' vs '{b_color}')")
        else:
            status = "🟢 Maru"
            reasons.append("Complete pair matched and verified successfully.")

        if "error" in ocr:
            reasons.append(f"[OCR error: {ocr['error']}]")

        display_top = f"`{t_text}`" if t_text else "*Blank*"
        display_bottom = f"`{b_text}`" if b_text else "*Blank*"

        report += (
            f"| {col_num} | {display_top} | {display_bottom} | **{t_color}** | **{b_color}** "
            f"| {status} | {', '.join(reasons)} |\n"
        )

    return ui_display_image, report


# ---------- UI CSS ----------
# ============================================================
# GRADIO UI CSS (FIXED FOR READABLE TABLES)
# ============================================================

apple_dark_pink_css = """
body, .gradio-container {
    background-color: #0f1115 !important;
}

h1, h2, h3, p, label {
    color: #ffffff !important;
}

button.primary {
    background: #f472b6 !important;
    color: #000000 !important;
    border: none !important;
    font-weight: bold !important;
}

/* --- FIX FOR REPORT TABLE VISIBILITY --- */
.prose, .prose * {
    color: #ffffff !important; /* Forces all text in Markdown/Report to solid white */
}

.prose table {
    background-color: #1a1d24 !important; /* Adds contrast behind the table */
    border-collapse: collapse !important;
    width: 100% !important;
}

.prose th {
    background-color: #272b35 !important;
    color: #f472b6 !important; /* Highlights headers in theme accent */
    border: 1px solid #3f4452 !important;
    padding: 8px !important;
}

.prose td {
    color: #ffffff !important; /* Clear white text inside table cells */
    border: 1px solid #2f3441 !important;
    padding: 8px !important;
}

.prose code {
    background-color: #2e3440 !important;
    color: #a3be8c !important; /* Greenish glow for OCR tag text blocks */
    padding: 2px 6px !important;
    border-radius: 4px !important;
}

footer {
    display: none !important;
}
"""
with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="pink"),
    css=apple_dark_pink_css
) as demo:
    gr.Markdown(
"# AI-Based Visual Inspection of Wire Terminal Connections"
)


    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="numpy", label="Upload Terminal Image")
            submit_btn = gr.Button("Process & Verify Structure", variant="primary")

        with gr.Column():
            output_img = gr.Image(type="numpy", label="Segmentation View Matrix")

    gr.Markdown("---")
    output_report = gr.Markdown(label="Verification Report Matrix")

    submit_btn.click(
        fn=process_and_verify,
        inputs=input_img,
        outputs=[output_img, output_report]
    )

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