import cv2
import numpy as np
import gradio as gr
from PIL import Image
import tensorflow as tf
import keras
from huggingface_hub import snapshot_download
import pytesseract
import io
import math
import os
import re
import tempfile
from collections import defaultdict
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from reportlab.lib.pagesizes import letter
from reportlab.lib.units import inch
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image as RLImage

import requests

_msi_model = None
LOGO_PATH = "logo.png"
COUNTAPI_BASE = "https://countapi.mileshilliard.com/api/v1"
COUNTAPI_KEY = "haanilango-design-analyzer-images-analyzed"


def format_count_badge(count):
    return (
        f'<div style="background-color:#1e2761; color:white; text-align:center; '
        f'padding:20px; border-radius:8px; margin-bottom:10px;">'
        f'<div style="font-size:20px; font-weight:bold;">{count:,} images were tested</div>'
        f'</div>'
    )


def get_current_count():
    try:
        resp = requests.get(f"{COUNTAPI_BASE}/get/{COUNTAPI_KEY}", timeout=5)
        if resp.status_code == 200:
            return resp.json().get("value", 0)
    except Exception:
        pass
    return 0


def increment_count():
    try:
        resp = requests.get(f"{COUNTAPI_BASE}/hit/{COUNTAPI_KEY}", timeout=5)
        if resp.status_code == 200:
            return resp.json().get("value", None)
    except Exception:
        pass
    return None


def add_watermark(pil_img, logo_path=LOGO_PATH, scale=0.20, opacity=0.95):
    if not os.path.exists(logo_path):
        return pil_img
    try:
        base = pil_img.convert("RGBA")
        logo = Image.open(logo_path).convert("RGBA")
        bbox = logo.getbbox()
        if bbox:
            logo = logo.crop(bbox)
        img_w, img_h = base.size
        logo_w = max(50, int(img_w * scale))
        logo_ratio = logo.height / logo.width
        logo_h = int(logo_w * logo_ratio)
        logo_resized = logo.resize((logo_w, logo_h), Image.LANCZOS)

        r, g, b, a = logo_resized.split()
        a = a.point(lambda p: int(p * opacity))
        logo_resized = Image.merge("RGBA", (r, g, b, a))

        margin = max(10, int(img_w * 0.015))
        pos = (img_w - logo_w - margin, img_h - logo_h - margin)

        base.paste(logo_resized, pos, logo_resized)
        return base.convert("RGB")
    except Exception:
        return pil_img


def load_msi_net():
    global _msi_model
    if _msi_model is None:
        hf_dir = snapshot_download(repo_id="alexanderkroner/MSI-Net")
        _msi_model = keras.layers.TFSMLayer(hf_dir, call_endpoint="serving_default")
    return _msi_model


def get_target_shape(original_shape):
    ar = original_shape[0] / original_shape[1]
    square_mode = abs(ar - 1.0)
    landscape_mode = abs(ar - 240 / 320)
    portrait_mode = abs(ar - 320 / 240)
    best = min(square_mode, landscape_mode, portrait_mode)
    if best == square_mode:
        return (320, 320)
    elif best == landscape_mode:
        return (240, 320)
    else:
        return (320, 240)


def preprocess_input(input_image, target_shape):
    t = tf.expand_dims(input_image, axis=0)
    t = tf.image.resize(t, target_shape, preserve_aspect_ratio=True)
    vp = target_shape[0] - t.shape[1]
    hp = target_shape[1] - t.shape[2]
    v1, v2 = vp // 2, vp - vp // 2
    h1, h2 = hp // 2, hp - hp // 2
    t = tf.pad(t, [[0, 0], [v1, v2], [h1, h2], [0, 0]])
    return t, [v1, v2], [h1, h2]


def postprocess_output(output_tensor, vp, hp, original_shape):
    output_tensor = output_tensor[:, vp[0]:output_tensor.shape[1] - vp[1], hp[0]:output_tensor.shape[2] - hp[1], :]
    output_tensor = tf.image.resize(output_tensor, original_shape)
    return output_tensor.numpy().squeeze()


def compute_text_boost_map(pil_img):
    img_arr = np.array(pil_img.convert("RGB"))
    h, w = img_arr.shape[:2]
    data = pytesseract.image_to_data(img_arr, config="--psm 11", output_type=pytesseract.Output.DICT)
    text_map = np.zeros((h, w), dtype=np.float32)
    for i in range(len(data['text'])):
        word = data['text'][i].strip()
        tw_, th_ = data['width'][i], data['height'][i]
        if not word or th_ < 8 or tw_ < 5:
            continue
        x, y = data['left'][i], data['top'][i]
        weight = min(1.0, (th_ / h) * 6)
        pad = int(th_ * 0.15)
        y0, y1 = max(0, y - pad), min(h, y + th_ + pad)
        x0, x1 = max(0, x - pad), min(w, x + tw_ + pad)
        text_map[y0:y1, x0:x1] = np.maximum(text_map[y0:y1, x0:x1], weight)
    return text_map


def compute_saliency_map(pil_image):
    model = load_msi_net()
    input_image = np.array(pil_image.convert("RGB"), dtype=np.float32)
    original_shape = input_image.shape[:2]
    target_shape = get_target_shape(original_shape)
    input_tensor, v_pad, h_pad = preprocess_input(input_image, target_shape)
    raw_output = model(input_tensor)
    output_tensor = list(raw_output.values())[0] if isinstance(raw_output, dict) else raw_output
    msi_saliency = postprocess_output(output_tensor, v_pad, h_pad, original_shape)
    msi_saliency = msi_saliency.astype(np.float32)
    msi_saliency -= msi_saliency.min()
    if msi_saliency.max() > 0:
        msi_saliency /= msi_saliency.max()

    text_boost = compute_text_boost_map(pil_image)
    combined = 0.7 * msi_saliency + 0.3 * text_boost
    return combined


def render_heatmap_overlay(img_bgr, saliency_map, alpha=0.45):
    heat_u8 = (saliency_map * 255).astype(np.uint8)
    heat_color = cv2.applyColorMap(heat_u8, cv2.COLORMAP_JET)
    return cv2.addWeighted(heat_color, alpha, img_bgr, 1 - alpha, 0)


def compute_attention_scores(saliency_map):
    flat = np.sort(saliency_map.flatten())
    n = len(flat)
    cum = np.cumsum(flat)
    gini = (n + 1 - 2 * np.sum(cum) / (cum[-1] + 1e-8)) / n
    focus_score = float(gini) * 100
    spread_score = float((saliency_map >= 0.5).mean()) * 100
    return {"focus_score": round(focus_score, 1), "spread_score": round(spread_score, 1)}


def compute_readability_score(read):
    if read['words_checked'] == 0:
        return None
    penalty = 12 * math.sqrt(read['contrast_issues']) + 8 * math.sqrt(read['small_text_issues'])
    return round(max(0, 100 - penalty), 1)


def coverage_grade_score(spread_score):
    if spread_score < 40:
        return 100.0
    return max(0.0, 100.0 - (spread_score - 40) * 2)


def compute_overall_grade(focus_score, spread_score, readability_score):
    coverage_score = coverage_grade_score(spread_score)
    if readability_score is None:
        total = round((0.55 / 0.75) * focus_score + (0.20 / 0.75) * coverage_score)
    else:
        total = round(0.55 * focus_score + 0.20 * coverage_score + 0.25 * readability_score)
    total = max(0, min(100, total))
    if total >= 75:
        grade = "Strong"
    elif total >= 50:
        grade = "Good Start"
    else:
        grade = "Room to Grow"
    return total, grade


def explain_grade(focus_score, spread_score, readability_score, read):
    coverage_score = coverage_grade_score(spread_score)
    weighted_coverage_deficit = 0.20 * max(0, 100 - coverage_score)
    if readability_score is None:
        if weighted_coverage_deficit > 0.55 * max(0, 100 - focus_score):
            return "based on attention alone (this design's text couldn't be reliably read, so readability wasn't scored) - mainly because attention is spread too widely without settling anywhere"
        elif focus_score < 50:
            return "based on attention alone (this design's text couldn't be reliably read, so readability wasn't scored) - mainly because attention is spread out without one clear focal point"
        elif focus_score > 90:
            return "based on attention alone (this design's text couldn't be reliably read, so readability wasn't scored) - attention is very tightly focused on one spot"
        else:
            return "based on attention alone - this design's text couldn't be reliably read, so readability wasn't included"
    total_issues = read['contrast_issues'] + read['small_text_issues']
    weighted_focus_deficit = 0.55 * max(0, 100 - focus_score)
    weighted_read_deficit = 0.25 * max(0, 100 - readability_score)

    if weighted_coverage_deficit > weighted_focus_deficit and weighted_coverage_deficit > weighted_read_deficit:
        return "mainly because attention is spread too widely across the design without settling anywhere"
    elif weighted_read_deficit >= weighted_focus_deficit and total_issues > 0:
        return "mainly due to text readability - some text is harder to read than ideal"
    elif focus_score < 50:
        return "mainly because attention is spread out, without one clear focal point"
    elif focus_score > 90:
        return "attention is very tightly focused on one spot - worth checking other key elements (logo, CTA) aren't being missed"
    else:
        return "a mix of small readability and focus factors - see details below"


def compute_hierarchy(text_sizes):
    if len(text_sizes) < 2:
        return None
    largest = max(text_sizes)
    smallest = min(text_sizes)
    ratio = largest / max(smallest, 0.1)
    score = round(min(100, ratio * 20), 1)
    return score, ratio


def hierarchy_description(ratio):
    if ratio < 1.5:
        return "text sizes are very similar - there's no single element that clearly leads the eye"
    elif ratio < 2.5:
        return "there's some size variation, but the hierarchy could be stronger"
    else:
        return "clear size hierarchy - one element leads, the rest support it"


def compute_balance(saliency_map):
    h, w = saliency_map.shape
    left = saliency_map[:, :w // 2].mean()
    right = saliency_map[:, w // 2:].mean()
    top = saliency_map[:h // 2, :].mean()
    bottom = saliency_map[h // 2:, :].mean()
    horiz_diff = abs(left - right) / max(left + right, 1e-6)
    vert_diff = abs(top - bottom) / max(top + bottom, 1e-6)
    imbalance = (horiz_diff + vert_diff) / 2
    score = round(max(0, 100 - imbalance * 200), 1)
    return score


def balance_description(score):
    if score >= 75:
        return "well balanced - visual weight is evenly distributed"
    elif score >= 50:
        return "somewhat uneven - one side carries noticeably more visual weight"
    else:
        return "heavily lopsided - attention is pulled hard toward one side or corner"


def detect_price_info(all_words):
    joined = " ".join(all_words)
    price_pattern = r'[\$£€¥]\s?\d[\d,]*(\.\d{1,2})?|\b\d[\d,]*(\.\d{1,2})?\s?(SGD|USD|GBP|EUR|dollars?|cents?)\b'
    return bool(re.search(price_pattern, joined, re.IGNORECASE))


def detect_location_info(all_words):
    joined = " ".join(all_words)
    location_pattern = (
        r'\b(venue|address|location|street|st\.|road|rd\.|avenue|ave\.|'
        r'blvd|drive|dr\.|www\.|\.com|\.sg|\.net)\b'
    )
    postal_pattern = r'\bS?\d{6}\b'
    return bool(re.search(location_pattern, joined, re.IGNORECASE)) or bool(re.search(postal_pattern, joined))


def coverage_description(spread_score):
    if spread_score < 15:
        return "concentrated on one clear focal point - that's usually a good sign, not a problem"
    elif spread_score < 40:
        return "landing mostly in a couple of areas, which is fairly typical"
    else:
        return "spread fairly widely across the design"


def create_score_chart(focus_score, spread_score, readability_score, hierarchy_score=None, balance_score=None):
    labels = ["Clear Focal\nPoint", "Attention\nCoverage"]
    values = [focus_score, spread_score]
    if readability_score is not None:
        labels.append("Readability")
        values.append(readability_score)
    if hierarchy_score is not None:
        labels.append("Visual\nHierarchy")
        values.append(hierarchy_score)
    if balance_score is not None:
        labels.append("Balance")
        values.append(balance_score)
    colors = []
    for v in values:
        if v >= 75:
            colors.append("#4CAF50")
        elif v >= 50:
            colors.append("#FFC107")
        else:
            colors.append("#F44336")
    fig, ax = plt.subplots(figsize=(5, 3), dpi=100)
    bars = ax.barh(labels, values, color=colors)
    ax.set_xlim(0, 100)
    ax.set_xlabel("Score out of 100")
    ax.invert_yaxis()
    for bar, v in zip(bars, values):
        ax.text(min(v + 2, 92), bar.get_y() + bar.get_height() / 2, f"{v:.0f}", va="center", fontsize=10)
    ax.set_title("Scores at a Glance")
    fig.tight_layout()
    buf = io.BytesIO()
    fig.savefig(buf, format="png")
    plt.close(fig)
    buf.seek(0)
    return Image.open(buf).convert("RGB")


def generate_recommendations(saliency_map, scores, read, img_h, img_w):
    tips = []
    bottom_band = saliency_map[int(img_h * 0.85):, :]
    bottom_avg = bottom_band.mean()
    if bottom_avg < 0.25:
        tips.append("Your bottom section (often where CTAs or contact info sit) is getting low attention. Consider bolder color contrast or larger text there.")

    if scores['focus_score'] < 40:
        tips.append("Attention is scattered with no clear focal point. Consider making one element (headline, product, or offer) visually dominant.")
    elif scores['focus_score'] > 90:
        tips.append("Attention is very narrowly focused on one spot - double check other important elements (logo, CTA) are not being ignored.")

    if read['contrast_issues'] > 0:
        if read['contrast_issues'] > 5:
            tips.append(f"A quick win: {read['contrast_issues']} text elements could read more clearly. The good news - the pattern (see white boxes marked \"Low contrast\" on the image) suggests one color choice affects most of them, so a single tweak will likely fix most at once.")
        else:
            tips.append(f"{read['contrast_issues']} text element(s) could be easier to read - see the white boxes marked \"Low contrast\" on the image above, then darken the text or lighten its background.")
    if read['small_text_issues'] > 0:
        if read['small_text_issues'] > 5:
            tips.append(f"Another quick win: {read['small_text_issues']} text elements are a bit small. Likely one font-size setting affects most of them, so this is usually a fast fix.")
        else:
            tips.append(f"{read['small_text_issues']} text element(s) are a bit small - see the white boxes marked \"Small text\" on the image above.")

    if not tips:
        tips.append("No major issues detected - this design is in solid shape.")

    return tips


def identify_hotspots(saliency_map, img_h, img_w, top_n=4):
    threshold = np.percentile(saliency_map, 85)
    binary = (saliency_map >= threshold).astype(np.uint8)
    kernel = np.ones((25, 25), np.uint8)
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(binary, connectivity=8)
    zones = []
    for i in range(1, num_labels):
        area = stats[i, cv2.CC_STAT_AREA]
        if area < (img_h * img_w * 0.005):
            continue
        cx, cy = centroids[i]
        h_pos = "top" if cy < img_h * 0.33 else ("bottom" if cy > img_h * 0.66 else "middle")
        w_pos = "left" if cx < img_w * 0.33 else ("right" if cx > img_w * 0.66 else "center")
        avg_intensity = float(saliency_map[labels == i].mean())
        zones.append({"position": f"{h_pos}-{w_pos}", "area_pct": round(float(area / (img_h * img_w)) * 100, 1),
                       "intensity": round(avg_intensity * 100, 1), "cx": int(cx), "cy": int(cy)})
    zones.sort(key=lambda z: -z["intensity"])
    top_zones = zones[:top_n]
    for idx, z in enumerate(top_zones):
        z["intensity_rank"] = idx + 1
    top_zones.sort(key=lambda z: (z["cy"], z["cx"]))
    return top_zones


def relative_luminance(rgb):
    def chan(c):
        c = c / 255.0
        return c / 12.92 if c <= 0.03928 else ((c + 0.055) / 1.055) ** 2.4
    r, g, b = rgb
    return 0.2126 * chan(r) + 0.7152 * chan(g) + 0.0722 * chan(b)


def contrast_ratio(rgb1, rgb2):
    l1 = relative_luminance(rgb1)
    l2 = relative_luminance(rgb2)
    lighter, darker = max(l1, l2), min(l1, l2)
    return (lighter + 0.05) / (darker + 0.05)


def required_contrast_ratio(size_pct):
    return 3.0 if size_pct >= 3.0 else 4.5


def contrast_severity_label(ratio, required_ratio):
    deficit = (required_ratio - ratio) / required_ratio
    if deficit > 0.6:
        return "could be much clearer"
    elif deficit > 0.35:
        return "could be clearer"
    elif deficit > 0.15:
        return "slightly less clear than ideal"
    else:
        return "close to the recommended level"


def sample_text_and_bg_color(img_arr, x, y, w, h):
    pad = max(2, int(h * 0.3))
    y0, y1 = max(0, y - pad), min(img_arr.shape[0], y + h + pad)
    x0, x1 = max(0, x - pad), min(img_arr.shape[1], x + w + pad)
    region = img_arr[y0:y1, x0:x1].reshape(-1, 3).astype(np.float32)
    if len(region) < 10:
        return None, None
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
    try:
        _, _, centers = cv2.kmeans(region, 2, None, criteria, 3, cv2.KMEANS_PP_CENTERS)
    except cv2.error:
        return None, None
    color1 = tuple(centers[0].astype(int))
    color2 = tuple(centers[1].astype(int))
    return color1, color2


def analyze_readability(pil_img):
    img_arr = np.array(pil_img.convert("RGB"))
    img_h = img_arr.shape[0]

    gray = cv2.cvtColor(img_arr, cv2.COLOR_RGB2GRAY)

    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    gray_enhanced = clahe.apply(gray)

    gray_denoised = cv2.medianBlur(gray_enhanced, 3)

    bw = cv2.adaptiveThreshold(gray_denoised, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 31, 10)

    cleanup_kernel = np.ones((2, 2), np.uint8)
    bw = cv2.morphologyEx(bw, cv2.MORPH_OPEN, cleanup_kernel)

    ocr_data = pytesseract.image_to_data(bw, output_type=pytesseract.Output.DICT)
    all_words_loose = [t.strip() for t in ocr_data['text'] if t.strip()]

    words_checked = 0
    contrast_issues = 0
    small_text_issues = 0
    issue_lines = []
    issue_boxes = []
    text_sizes = []
    for i in range(len(ocr_data['text'])):
        word = ocr_data['text'][i].strip()
        conf = int(ocr_data['conf'][i])
        if not word or conf < 60 or len(word) < 2:
            continue
        x, y, w, h = ocr_data['left'][i], ocr_data['top'][i], ocr_data['width'][i], ocr_data['height'][i]
        if w < 3 or h < 3:
            continue
        words_checked += 1
        size_pct = (h / img_h) * 100
        text_sizes.append(size_pct)
        text_color, bg_color = sample_text_and_bg_color(img_arr, x, y, w, h)
        if text_color is None:
            continue
        ratio = contrast_ratio(text_color, bg_color)
        required_ratio = required_contrast_ratio(size_pct)
        is_contrast_issue = ratio < required_ratio
        is_small_issue = size_pct < 1.2

        if is_contrast_issue:
            contrast_issues += 1
            if len(issue_lines) < 5:
                severity = contrast_severity_label(ratio, required_ratio)
                issue_lines.append(f"- \"{word}\" {severity} against its background")
        if is_small_issue:
            small_text_issues += 1
            if len(issue_lines) < 5:
                issue_lines.append(f"- \"{word}\" is a bit small to read comfortably")

        if is_contrast_issue or is_small_issue:
            combined_label = " + ".join(
                l for l, flag in [("Low contrast", is_contrast_issue), ("Small text", is_small_issue)] if flag
            )
            issue_boxes.append({"box": (x, y, w, h), "label": combined_label})
    return {"words_checked": words_checked, "contrast_issues": contrast_issues,
            "small_text_issues": small_text_issues, "issue_lines": issue_lines, "issue_boxes": issue_boxes,
            "text_sizes": text_sizes, "all_words": all_words_loose}


def merge_issue_boxes(issue_boxes, gap=15):
    by_label = defaultdict(list)
    for item in issue_boxes:
        by_label[item["label"]].append(item["box"])
    merged = []
    for label, boxes in by_label.items():
        boxes = sorted(boxes, key=lambda b: (b[1], b[0]))
        used = [False] * len(boxes)
        for i in range(len(boxes)):
            if used[i]:
                continue
            x, y, w, h = boxes[i]
            mx0, my0, mx1, my1 = x, y, x + w, y + h
            used[i] = True
            changed = True
            while changed:
                changed = False
                for j in range(len(boxes)):
                    if used[j]:
                        continue
                    bx, by, bw, bh = boxes[j]
                    bx0, by0, bx1, by1 = bx, by, bx + bw, by + bh
                    vertical_overlap = min(my1, by1) - max(my0, by0)
                    if vertical_overlap > 0 and bx0 - mx1 <= gap and bx1 >= mx0 - gap:
                        mx0, my0 = min(mx0, bx0), min(my0, by0)
                        mx1, my1 = max(mx1, bx1), max(my1, by1)
                        used[j] = True
                        changed = True
            merged.append({"box": (mx0, my0, mx1 - mx0, my1 - my0), "label": label})
    return merged


def draw_issue_markers(overlay_bgr, issue_boxes, img_h, img_w):
    img = overlay_bgr.copy()
    scale = max(1.0, min(img_h, img_w) / 800.0)
    merged_issues = merge_issue_boxes(issue_boxes)
    placed_label_rects = []
    pad = int(4 * scale)
    box_border = max(2, int(3 * scale))
    thin_border = max(1, int(scale))
    font_scale = 0.5 * scale
    text_thickness = max(1, int(round(scale)))
    for issue in merged_issues:
        x, y, w, h = issue["box"]
        cv2.rectangle(img, (x - pad, y - pad), (x + w + pad, y + h + pad), (255, 255, 255), box_border)
        cv2.rectangle(img, (x - pad, y - pad), (x + w + pad, y + h + pad), (0, 0, 0), thin_border)
        label = issue["label"]
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, text_thickness)
        label_x0 = x - pad
        label_x1 = label_x0 + tw + int(8 * scale)
        label_y1 = max(th + int(6 * scale), y - pad)
        label_y0 = label_y1 - th - int(6 * scale)

        collision = True
        attempts = 0
        while collision and attempts < 10:
            collision = False
            for (px0, py0, px1, py1) in placed_label_rects:
                if not (label_x1 < px0 or label_x0 > px1 or label_y1 < py0 or label_y0 > py1):
                    collision = True
                    label_y1 -= (th + int(10 * scale))
                    label_y0 = label_y1 - th - int(6 * scale)
                    break
            attempts += 1

        placed_label_rects.append((label_x0, label_y0, label_x1, label_y1))
        cv2.rectangle(img, (label_x0, label_y0), (label_x1, label_y1), (255, 255, 255), -1)
        cv2.rectangle(img, (label_x0, label_y0), (label_x1, label_y1), (0, 0, 0), thin_border)
        cv2.putText(img, label, (label_x0 + int(4 * scale), label_y1 - int(4 * scale)),
                    cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), text_thickness, cv2.LINE_AA)
    return img


def markdown_to_paragraphs(md_text, styles):
    story = []
    for raw_line in md_text.split("\n"):
        line = raw_line.strip()
        if not line:
            story.append(Spacer(1, 6))
            continue
        if line.startswith("## "):
            story.append(Paragraph(line[3:], styles['Heading1']))
        elif line.startswith("### "):
            story.append(Paragraph(line[4:], styles['Heading2']))
        elif line.startswith("---"):
            story.append(Spacer(1, 10))
        elif line.startswith("- ") or line.startswith("* "):
            text = line[2:]
            text = re.sub(r'\*\*(.*?)\*\*', r'<b>\1</b>', text)
            text = re.sub(r'\*(.*?)\*', r'<i>\1</i>', text)
            story.append(Paragraph("&bull; " + text, styles['Normal']))
        elif re.match(r'^\d+\.\s', line):
            num, text = line.split('.', 1)
            text = re.sub(r'\*\*(.*?)\*\*', r'<b>\1</b>', text.strip())
            story.append(Paragraph(f"{num}. {text}", styles['Normal']))
        else:
            text = re.sub(r'\*\*(.*?)\*\*', r'<b>\1</b>', line)
            text = re.sub(r'\*(.*?)\*', r'<i>\1</i>', text)
            text = re.sub(r'\[([^\]]+)\]\(([^)]+)\)', r'<link href="\2"><u>\1</u></link>', text)
            story.append(Paragraph(text, styles['Normal']))
    return story


def stamp_pdf_watermark(canvas_obj, doc):
    if not os.path.exists(LOGO_PATH):
        return
    try:
        from reportlab.lib.utils import ImageReader
        logo = Image.open(LOGO_PATH).convert("RGBA")
        bbox = logo.getbbox()
        if bbox:
            logo = logo.crop(bbox)
        logo_h = 0.8 * inch
        logo_w = logo_h * (logo.width / logo.height)
        margin = 15
        canvas_obj.saveState()
        canvas_obj.setFillAlpha(0.95)
        canvas_obj.drawImage(ImageReader(logo), doc.pagesize[0] - logo_w - margin, margin,
                              width=logo_w, height=logo_h, mask='auto', preserveAspectRatio=True)
        canvas_obj.restoreState()
    except Exception:
        pass


def export_pdf(heatmap_img, chart_img, summary_md):
    if heatmap_img is None or not summary_md:
        return None
    if isinstance(heatmap_img, np.ndarray):
        heatmap_img = Image.fromarray(heatmap_img)
    if chart_img is not None and isinstance(chart_img, np.ndarray):
        chart_img = Image.fromarray(chart_img)
    styles = getSampleStyleSheet()
    tmp_dir = tempfile.gettempdir()
    pdf_path = os.path.join(tmp_dir, "design_analysis_report.pdf")
    doc = SimpleDocTemplate(pdf_path, pagesize=letter, topMargin=0.6 * inch, bottomMargin=0.6 * inch,
                             leftMargin=0.7 * inch, rightMargin=0.7 * inch)
    story = [Paragraph("Design Analyzer Report", styles['Title']), Spacer(1, 12)]

    max_width = 6.1 * inch
    heatmap_path = os.path.join(tmp_dir, "heatmap_export_temp.png")
    heatmap_img.save(heatmap_path)
    ratio = heatmap_img.height / heatmap_img.width
    story.append(RLImage(heatmap_path, width=max_width, height=max_width * ratio))
    story.append(Spacer(1, 14))

    if chart_img is not None:
        chart_path = os.path.join(tmp_dir, "chart_export_temp.png")
        chart_img.save(chart_path)
        chart_ratio = chart_img.height / chart_img.width
        story.append(RLImage(chart_path, width=max_width, height=max_width * chart_ratio))
        story.append(Spacer(1, 14))

    story.extend(markdown_to_paragraphs(summary_md, styles))
    doc.build(story, onFirstPage=stamp_pdf_watermark, onLaterPages=stamp_pdf_watermark)
    return pdf_path


def draw_scan_path(overlay_bgr, zones, img_h, img_w):
    img = overlay_bgr.copy()
    scale = max(1.0, min(img_h, img_w) / 800.0)
    outer_thickness = int(8 * scale)
    inner_thickness = int(4 * scale)
    ordered = sorted(zones, key=lambda z: z["intensity_rank"])
    for i in range(len(ordered) - 1):
        pt1 = (ordered[i]["cx"], ordered[i]["cy"])
        pt2 = (ordered[i + 1]["cx"], ordered[i + 1]["cy"])
        cv2.arrowedLine(img, pt1, pt2, (0, 0, 0), outer_thickness, tipLength=0.08, line_type=cv2.LINE_AA)
        cv2.arrowedLine(img, pt1, pt2, (0, 255, 255), inner_thickness, tipLength=0.08, line_type=cv2.LINE_AA)
    return img


def draw_zone_labels(overlay_bgr, zones, img_h, img_w):
    img = overlay_bgr.copy()
    scale = max(1.0, min(img_h, img_w) / 800.0)
    radius = int(22 * scale)
    font_scale = 0.9 * scale
    thickness = max(2, int(3 * scale))
    for i, z in enumerate(zones, 1):
        cx, cy = z["cx"], z["cy"]
        cv2.circle(img, (cx, cy), radius, (255, 255, 255), -1)
        cv2.circle(img, (cx, cy), radius, (0, 0, 0), thickness)
        text = str(i)
        (tw, th), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness)
        cv2.putText(img, text, (cx - tw // 2, cy + th // 2), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), thickness, cv2.LINE_AA)
    return img


def flatten_transparency(pil_image):
    if pil_image.mode in ("RGBA", "LA") or (pil_image.mode == "P" and "transparency" in pil_image.info):
        rgba = pil_image.convert("RGBA")
        background = Image.new("RGB", rgba.size, (255, 255, 255))
        background.paste(rgba, mask=rgba.split()[-1])
        return background
    return pil_image.convert("RGB")


def analyze(pil_image, overlay_strength):
    if pil_image is None:
        return None, None, "Upload an image first.", gr.update()
    pil_rgb = flatten_transparency(pil_image)
    img_bgr = cv2.cvtColor(np.array(pil_rgb), cv2.COLOR_RGB2BGR)
    h, w = img_bgr.shape[:2]

    saliency_map = compute_saliency_map(pil_rgb)
    overlay_bgr = render_heatmap_overlay(img_bgr, saliency_map, alpha=overlay_strength)

    scores = compute_attention_scores(saliency_map)
    zones = identify_hotspots(saliency_map, h, w)
    read = analyze_readability(pil_rgb)
    readability_score = compute_readability_score(read)
    total_score, grade = compute_overall_grade(scores['focus_score'], scores['spread_score'], readability_score)
    grade_reason = explain_grade(scores['focus_score'], scores['spread_score'], readability_score, read)
    tips = generate_recommendations(saliency_map, scores, read, h, w)
    hierarchy_result = compute_hierarchy(read.get('text_sizes', []))
    balance_score = compute_balance(saliency_map)
    hierarchy_score = hierarchy_result[0] if hierarchy_result else None
    chart_image = create_score_chart(scores['focus_score'], scores['spread_score'], readability_score,
                                      hierarchy_score, balance_score)

    overlay_bgr = draw_scan_path(overlay_bgr, zones, h, w)
    overlay_bgr = draw_zone_labels(overlay_bgr, zones, h, w)
    overlay_bgr = draw_issue_markers(overlay_bgr, read.get("issue_boxes", []), h, w)
    overlay_rgb = cv2.cvtColor(overlay_bgr, cv2.COLOR_BGR2RGB)

    summary = f"## Grade: {grade} — {total_score}/100\n"
    summary += f"*Why: {grade_reason}*\n\n"

    summary += "**Where Attention Goes First:** *(numbers match the image above)*\n"
    if zones:
        for i, z in enumerate(zones, 1):
            summary += f"{i}. {z['position'].replace('-', ' ').title()} ({z['intensity']}/100 intensity, {z['area_pct']}% of design)\n"
    else:
        summary += "- Attention is fairly even across the design, no strong single hotspot.\n"

    summary += (
        f"\n**Scores:**\n"
        f"- Clear Focal Point: {scores['focus_score']}/100 (how strongly attention lands on one main spot, vs scattered everywhere)\n"
        f"- Attention Coverage: {scores['spread_score']}/100 - {coverage_description(scores['spread_score'])}\n"
        f"- Readability: "
    )
    if read['words_checked'] == 0:
        summary += "couldn't be reliably read on this design (common with heavily stylized or hand-lettered fonts) - not included in the score above\n"
    elif read['contrast_issues'] == 0 and read['small_text_issues'] == 0:
        summary += "clean, no changes needed\n"
    else:
        total_issues = read['contrast_issues'] + read['small_text_issues']
        summary += f"{total_issues} easy improvement(s) spotted - see below\n"

    if read['issue_lines']:
        summary += "\n**A few examples:**\n" + "\n".join(read['issue_lines']) + "\n"
        summary += (
            "\n*Note: this is automated and can occasionally flag text that's actually fine, "
            "especially on busy or gradient backgrounds - if something flagged looks perfectly "
            "readable to your eye, trust your eye.*\n"
        )

    summary += "\n**How to Improve:** *(common and usually quick to fix)*\n"
    for tip in tips:
        summary += f"- {tip}\n"

    summary += "\n**Layout & Composition:** *(based on classic layout principles - hierarchy and balance)*\n"
    if hierarchy_result:
        summary += f"- Visual Hierarchy: {hierarchy_description(hierarchy_result[1])}\n"
    else:
        summary += "- Visual Hierarchy: not enough distinct text elements to judge\n"
    summary += f"- Balance: {balance_description(balance_score)}\n"

    has_price = detect_price_info(read.get('all_words', []))
    has_location = detect_location_info(read.get('all_words', []))
    summary += (
        "\n**Content Checklist:** *(Mike Stevens' classic Who/What/Where/Why/How Much - "
        "not every design needs all five)*\n"
        f"- How Much (price): {'found' if has_price else 'not detected - if this design should show a price, double check it'}\n"
        f"- Where (location/contact): {'found' if has_location else 'not detected - if this design should point somewhere, double check it'}\n"
        "- Who / What / Why: these need a human read, not automation - ask yourself: is it "
        "clear who's offering this, exactly what they're offering, and why someone should care?\n"
    )

    summary += (
        "\n---\n"
        "**Methodology:** Attention prediction combines MSI-Net, a peer-reviewed saliency model "
        "(Kroner et al., *Neural Networks*, 2020, "
        "[DOI: 10.1016/j.neunet.2020.05.004](https://doi.org/10.1016/j.neunet.2020.05.004)), "
        "with a size-proportional boost for legible text - since eye-tracking research confirms "
        "large headline text reliably draws attention, which pure bottom-up saliency models can "
        "underweight relative to small high-contrast graphics. This is a data-informed prediction, "
        "not a substitute for live user testing."
    )

    summary += (
        "\n\n---\n"
        "**This tool flags what's off. A designer knows what's right.** "
        "Haan is a brand identity and visual designer with 12+ years turning "
        "insights like these into designs that actually convert, stay true to "
        "your brand, and look intentional (not just \"fixed\").\n\n"
        "[Visit haanilango.com](https://www.haanilango.com) · Tel: +65 9027 2070"
    )

    result_image = add_watermark(Image.fromarray(overlay_rgb))
    chart_image = add_watermark(chart_image, scale=0.26)

    new_count = increment_count()
    badge_html = format_count_badge(new_count) if new_count is not None else gr.update()

    return result_image, chart_image, summary, badge_html


with gr.Blocks(title="Design Analyzer") as demo:
    count_display = gr.HTML(format_count_badge(0))
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload design")
            strength_slider = gr.Slider(0.1, 0.8, value=0.45, step=0.05, label="Heatmap overlay strength")
            analyze_btn = gr.Button("Analyze", variant="primary")
        with gr.Column():
            image_output = gr.Image(label="Heatmap result", interactive=False)
            chart_output = gr.Image(label="Scores at a glance", interactive=False)
    with gr.Row():
        text_output = gr.Markdown()
    with gr.Row():
        pdf_btn = gr.Button("Download PDF Report")
        pdf_output = gr.File(label="PDF Report")
    analyze_btn.click(fn=analyze, inputs=[image_input, strength_slider],
                       outputs=[image_output, chart_output, text_output, count_display])
    pdf_btn.click(fn=export_pdf, inputs=[image_output, chart_output, text_output], outputs=pdf_output)
    demo.load(fn=lambda: format_count_badge(get_current_count()), outputs=count_display)
    gr.Markdown("---\n© 2026 Haan Ilango. All rights reserved. This tool's methodology and design are original work.")

demo.launch()