"""Display-only projections. Raw model output is never rewritten."""
import base64
import html
import io
import json
import re

import bleach
from bs4 import BeautifulSoup
from markdown_it import MarkdownIt
from PIL import ImageDraw, ImageFont

from catalog import ROOT
from presentation import TASK_VIEWS, JSON_TASKS
from structured_json import strict_json_loads

COLORS = ["#0d9488", "#e05a47", "#7659b4", "#b88712", "#2876ba"]
RESULT_CSS = (ROOT / "assets/result.css").read_text().replace("__ASSETS__", str(ROOT / "assets"))
CROP = re.compile(r"(?:page_(\d+)/)?images/bbox_(\d+)_(\d+)_(\d+)_(\d+)\.jpg")
TAGS = {"p", "div", "span", "br", "hr", "h1", "h2", "h3", "h4", "h5", "h6", "b", "strong", "i", "em", "u", "s", "sub", "sup", "ul", "ol", "li", "blockquote", "pre", "code", "table", "thead", "tbody", "tfoot", "tr", "th", "td", "img"}


def parse_json(raw):
    text = raw.strip()
    fence = re.fullmatch(r"```(?:json)?\s*\n(.*?)\n```", text, re.S | re.I)
    return strict_json_loads(fence[1] if fence else text)


def output_structure_error(task, data):
    arrays = {"document_structured", "document_layout", "figure_captioning", "word_spotting",
              "reference_grounding", "crosspage_table"}
    expected = (dict, list) if task == "kie_schema" else (list,) if task in arrays else (dict,)
    if not isinstance(data, expected):
        name = "an object or array" if len(expected) == 2 else "an array" if list in expected else "an object"
        return f"Expected a JSON {name.removeprefix('an ')} for this task."
    if task in arrays and not all(isinstance(item, dict) for item in data):
        return "Expected a JSON array of objects for this task."
    return None


def bbox(value):
    if isinstance(value, str):
        match = re.fullmatch(r"\s*[\[(]\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*[\])]\s*", value)
        if not match:
            raise ValueError("Invalid bounding box string.")
        value = list(map(int, match.groups()))
    if not isinstance(value, (list, tuple)) or len(value) != 4 or not all(type(v) is int for v in value):
        raise ValueError("Invalid bounding box.")
    x1, y1, x2, y2 = value
    if not (0 <= x1 < x2 <= 1000 and 0 <= y1 < y2 <= 1000):
        raise ValueError("Bounding box is out of range or empty.")
    return value


def crop_image(image, value):
    x1, y1, x2, y2 = bbox(value)
    w, h = image.size
    return image.crop((round(x1*w/1000), round(y1*h/1000), round(x2*w/1000), round(y2*h/1000)))


def image_uri(image):
    output = io.BytesIO()
    image.convert("RGB").save(output, format="JPEG", quality=88)
    return "data:image/jpeg;base64," + base64.b64encode(output.getvalue()).decode("ascii")


def boxed_answer_rule(state, silent):
    if not state.src.startswith(r"\boxed", state.pos):
        return False
    opening = state.pos + len(r"\boxed")
    while opening < state.posMax and state.src[opening].isspace():
        opening += 1
    if opening >= state.posMax or state.src[opening] != "{":
        return False
    depth, end = 1, opening + 1
    # Balance nested TeX groups without consuming escaped braces or code spans.
    while end < state.posMax:
        char = state.src[end]
        if char == "\\":
            end += 2
            continue
        if char == "{":
            depth += 1
        elif char == "}":
            depth -= 1
            if depth == 0:
                if not silent:
                    token = state.push("boxed_answer", "", 0)
                    token.content = state.src[opening + 1:end]
                state.pos = end + 1
                return True
        end += 1
    return False


def boxed_answer_html(tokens, index, options, env):
    content = tokens[index].content
    # Plain-language answers retain spaces and upright text; TeX remains TeX.
    if re.search(r"[\\{}_^]", content):
        content = r"\(" + content + r"\)"
    return '<span class="boxed-answer">' + html.escape(content) + '</span>'


def markdown_html(text, *, boxed_answers=False):
    # Protect math before Markdown consumes backslashes and underscores.
    expressions = []
    def protect(match):
        expressions.append(match[0])
        return f"OVISMATHPLACEHOLDER{len(expressions)-1}END"
    text = re.sub(r"\$\$[\s\S]+?\$\$|\\\[[\s\S]+?\\\]|\\\([\s\S]+?\\\)|(?<![\\\w])\$(?!\s|\$)[^$\n]+?(?<!\s)\$(?!\w)", protect, text)
    markdown = MarkdownIt("commonmark", {"html": True, "breaks": True}).enable("table")
    if boxed_answers:
        markdown.inline.ruler.before("escape", "boxed_answer", boxed_answer_rule)
        markdown.renderer.rules["boxed_answer"] = boxed_answer_html
    rendered = markdown.render(text)
    for index, expression in enumerate(expressions):
        rendered = rendered.replace(f"OVISMATHPLACEHOLDER{index}END", html.escape(expression))
    return rendered


def sanitize(content, images):
    soup = BeautifulSoup(content, "html.parser")
    for element in soup.find_all(["script", "style", "iframe", "object", "embed", "link", "meta", "svg", "math"]):
        element.decompose()
    for image in soup.find_all("img"):
        match = CROP.fullmatch(str(image.get("src", "")))
        if not match:
            image.replace_with("[Unsupported image reference]")
            continue
        page = int(match[1] or 1) - 1
        try:
            if not 0 <= page < len(images):
                raise ValueError("Page out of range")
            image["src"] = image_uri(crop_image(images[page], list(map(int, match.groups()[1:]))))
            image.attrs = {"src": image["src"], "alt": "Document region"}
        except (ValueError, IndexError):
            image.replace_with("[Invalid image region]")
    return bleach.clean(str(soup), tags=TAGS,
                        attributes={"*": ["class"], "img": ["src", "alt"], "td": ["rowspan", "colspan"], "th": ["rowspan", "colspan"]},
                        protocols={"data"}, strip=True)


def display_html(content, images=()):
    safe = sanitize(content, images)
    # Give wide tables their own scroll area instead of breaking every word.
    soup = BeautifulSoup(safe, "html.parser")
    for table in soup.find_all("table"):
        if table.find_parent("table") is None:
            wrapper = soup.new_tag("div", attrs={"class": "table-scroll", "tabindex": "0", "role": "region", "aria-label": "Result table"})
            table.wrap(wrapper)
    return str(soup)


def needs_math(content):
    return bool(re.search(r'\$|\\\(|\\\[', content))


def frame_document(content, images=(), math=None):
    safe = display_html(content, images)
    if math is None:
        math = needs_math(safe)
    mathjax_path = f"./gradio_api/file={ROOT / 'assets' / 'mathjax.js'}"
    document = '''<!doctype html><html><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
<style>body{margin:0;padding:24px;color:#242a30;background:white;font:15px/1.65 system-ui,sans-serif;overflow-wrap:anywhere;letter-spacing:0}h1,h2,h3{line-height:1.3;font-size:20px;margin:18px 0 10px}h2{font-size:18px}h3{font-size:16px}p{margin:8px 0 14px}img{max-width:100%;height:auto}table{border-collapse:collapse;max-width:100%;display:block;overflow:auto;margin:14px 0;font-size:13px}td,th{border:1px solid #ccd5d7;padding:7px 10px;min-width:35px;vertical-align:top}th{background:#f0f5f4}pre{white-space:pre-wrap;font:13px/1.6 ui-monospace,monospace;background:#f4f6f6;padding:14px;overflow-wrap:anywhere}blockquote{margin-left:0;padding-left:14px;border-left:3px solid #0d9488}mjx-container[display="true"]{overflow-x:auto;overflow-y:hidden;padding:6px 0}@media(max-width:500px){body{padding:14px;font-size:14px}}</style>
<script>window.MathJax={tex:{inlineMath:[['$','$'],['\\\\(','\\\\)']],displayMath:[['$$','$$'],['\\\\[','\\\\]']]},svg:{fontCache:'local'},options:{enableMenu:false}};</script>
<script defer src="''' + html.escape(mathjax_path, quote=True) + '''"></script></head><body>''' + safe + "</body></html>"
    document = document.replace("</head>", "<style>" + RESULT_CSS + "</style></head>")
    if not math:
        document = re.sub(r"<script[\s\S]*?</script>", "", document)
    return document


def frame(content, images=(), math=None):
    document = frame_document(content, images, math)
    return '<iframe title="Rendered model output" sandbox="allow-scripts" referrerpolicy="no-referrer" style="width:100%;height:570px;border:0;background:white" srcdoc="' + html.escape(document, quote=True) + '"></iframe>'


def pending_output(task, raw="", state="Ready"):
    title = TASK_VIEWS[task][1]
    if raw:
        return frame('<div class="result-meta">' + html.escape(state) + '</div><pre class="stream">' + html.escape(raw) + '</pre>', math=False)
    label = "No result yet" if state == "Ready" else state
    return frame(f'<div class="empty-result"><div class="empty-glyph"></div><h3>{html.escape(title)}</h3><p>{html.escape(label)}</p></div>', math=False)


def value_html(value):
    if isinstance(value, dict):
        rows = []
        for key, item in value.items():
            rows.append('<tr><th>' + html.escape(str(key)) + '</th><td>' + value_html(item) + '</td></tr>')
        return '<table class="field-table"><tbody>' + ''.join(rows) + '</tbody></table>'
    if isinstance(value, list):
        if not value:
            return '<span class="empty-value">[]</span>'
        if all(isinstance(item, dict) for item in value):
            keys = list(dict.fromkeys(key for item in value for key in item))
            header = '<tr>' + ''.join('<th>' + html.escape(str(key)) + '</th>' for key in keys) + '</tr>'
            rows = ''.join('<tr>' + ''.join('<td>' + (value_html(item[key]) if key in item else '<span class="empty-value">Not returned</span>') + '</td>' for key in keys) + '</tr>' for item in value)
            return '<table class="records-table"><thead>' + header + '</thead><tbody>' + rows + '</tbody></table>'
        return '<ol>' + ''.join('<li>' + value_html(item) + '</li>' for item in value) + '</ol>'
    text = json.dumps(value, ensure_ascii=False) if value is None or isinstance(value, bool) else str(value)
    return '<span class="field-value">' + html.escape(text) + '</span>'


def region_rows(data):
    records = data if isinstance(data, list) else [data]
    rows = []
    for index, record in enumerate(records):
        if not isinstance(record, dict):
            rows.append('<tr><td>' + str(index + 1) + '</td><td colspan="2">' + value_html(record) + '</td></tr>')
            continue
        text = record.get("text", record.get("answer", record.get("category", record.get("reference", ""))))
        coordinates = html.escape(str(record.get("bbox", "Not returned")))
        rows.append(f'<tr><td><span class="region-number">{index+1:02d}</span></td><td>{html.escape(str(text))}</td><td><code>{coordinates}</code></td></tr>')
    return '<table class="regions-table"><thead><tr><th>#</th><th>Content / category</th><th>Coordinates</th></tr></thead><tbody>' + ''.join(rows) + '</tbody></table>'


def annotated(image, data):
    result = image.copy()
    draw = ImageDraw.Draw(result)
    width = max(2, round(result.width / 500))
    try:
        font = ImageFont.truetype("DejaVuSans.ttf", max(16, result.width // 70))
    except OSError:
        font = ImageFont.load_default()
    records = data if isinstance(data, list) else [data]
    invalid = 0
    for index, record in enumerate(records):
        if not isinstance(record, dict) or "bbox" not in record:
            continue
        try:
            x1, y1, x2, y2 = bbox(record["bbox"])
        except ValueError:
            invalid += 1
            continue
        color = COLORS[index % len(COLORS)]
        x1, x2 = x1*result.width/1000, x2*result.width/1000
        y1, y2 = y1*result.height/1000, y2*result.height/1000
        draw.rectangle((x1, y1, x2, y2), outline=color, width=width)
        label = str(index + 1)
        box = draw.textbbox((0, 0), label, font=font)
        tw, th = box[2]+8, box[3]+5
        ly = max(0, y1-th)
        draw.rectangle((x1, ly, x1+tw, ly+th), fill=color)
        draw.text((x1+4, ly), label, fill="white", font=font)
    return result, invalid


def render_output(task, raw, images):
    data, overlays, warnings = None, [], []
    if task in JSON_TASKS:
        try:
            data = parse_json(raw)
        except (ValueError, TypeError) as exc:
            warnings.append(f"Invalid JSON: {exc}")
        else:
            structure_error = output_structure_error(task, data)
            if structure_error:
                warnings.append(structure_error)
                body = ('<div class="warning-panel"><strong>Unexpected output structure</strong>'
                        + html.escape(structure_error) + '</div><pre>' + html.escape(raw) + '</pre>')
                return frame(body, images), overlays, data, warnings
    if data is not None and task in {"document_structured", "document_layout", "figure_captioning", "word_spotting", "reference_grounding", "position_aware_vqa"}:
        image, invalid = annotated(images[0], data)
        overlays.append((image, "Detected regions"))
        if invalid:
            warnings.append(f"{invalid} invalid bounding box(es); retained in raw output.")
    meta = '<div class="result-meta"><strong>' + TASK_VIEWS[task][1] + '</strong><span>' + TASK_VIEWS[task][2] + '</span></div>'
    if data is not None and task in {"document_structured", "figure_captioning"} and isinstance(data, list):
        parts, context = [], []
        for index, item in enumerate(data):
            if not isinstance(item, dict):
                context.append(value_html(item))
                continue
            content = str(item.get("content", ""))
            crop = ""
            if "<caption>" in content and "bbox" in item:
                try:
                    coordinates = bbox(item["bbox"])
                    ref = "images/bbox_" + "_".join(map(str, coordinates)) + ".jpg"
                    crop = f'<img src="{ref}">'
                except ValueError:
                    pass
            is_caption = "<caption>" in content
            content = re.sub(r"<mermaid>([\s\S]*?)</mermaid>", lambda m: "\n```text\n" + m[1] + "\n```\n", content)
            content = content.replace("<caption>", "").replace("</caption>", "")
            if task == "figure_captioning" and is_caption:
                label = html.escape(str(item.get("category", "Figure")))
                parts.append(f'<div class="figure-item"><div class="figure-label">{index+1:02d} / {label}</div>{crop}{markdown_html(content)}</div>')
            elif task == "figure_captioning":
                context.append(markdown_html(content))
            else:
                parts.append('<div class="document-block">' + crop + markdown_html(content) + '</div>')
        body = meta + "\n".join(parts)
        if context:
            body += '<div class="figure-context"><div class="result-meta">Page context</div>' + '\n'.join(context) + '</div>'
    elif data is not None and task == "crosspage_table" and isinstance(data, list):
        parts = []
        for item in data:
            if not isinstance(item, dict):
                parts.append(value_html(item))
                continue
            pages = item.get("source_pages", [])
            badges = ''.join('<span class="table-source">Page ' + html.escape(str(page)) + '</span>' for page in pages) if isinstance(pages, list) else value_html(pages)
            parts.append('<div class="merged-table">' + badges + str(item.get("html", "")) + '</div>')
        body = meta + '\n'.join(parts)
    elif data is not None and task in {"kie_open", "kie_fields", "kie_schema"}:
        body = meta + value_html(data)
    elif data is not None and task == "chart_parsing" and isinstance(data, dict):
        title = html.escape(str(data.get("title", "Chart data")))
        body = meta + '<h2 class="result-title">' + title + '</h2>'
        if "values" in data:
            body += value_html(data["values"])
        # Show every returned field, including unrecognized chart schema fields.
        metadata = {key: value for key, value in data.items() if key not in {"values", "title"}}
        if metadata:
            body += '<div class="figure-context">' + value_html(metadata) + '</div>'
    elif data is not None and task == "position_aware_vqa" and isinstance(data, dict):
        body = meta + '<div class="answer">' + html.escape(str(data.get("answer", ""))) + '</div>' + region_rows(data)
    elif data is not None and task in {"document_layout", "word_spotting", "reference_grounding"}:
        body = meta + region_rows(data)
        if task == "document_layout" and overlays:
            images = [*images, overlays[0][0]]
            body = meta + f'<img class="layout-image" src="page_{len(images)}/images/bbox_0_0_1000_1000.jpg">' + region_rows(data)
    elif data is not None:
        body = meta + value_html(data)
    elif task in JSON_TASKS:
        body = '<div class="warning-panel"><strong>Structured output unavailable</strong>Invalid JSON. Original model response:</div><pre>' + html.escape(raw) + '</pre>'
    elif task == "formula_recognition":
        expression = raw.strip()
        if not any(token in expression for token in ("$", r"\[", r"\(")):
            expression = r"\[" + expression + r"\]"
        body = meta + '<div class="formula">' + markdown_html(expression) + '</div>'
    elif task == "image_classification":
        body = meta + '<span class="answer-label">' + html.escape(raw) + '</span>'
    elif task == "visual_question_answering":
        body = meta + '<div class="answer">' + markdown_html(raw, boxed_answers=True) + '</div>'
    elif task == "document_markdown":
        body = markdown_html(raw)
    else:
        body = meta + '<pre class="text-output">' + html.escape(raw) + '</pre>'
    return frame(body, images), overlays, data, warnings