from __future__ import annotations

import logging
import os
from pathlib import Path

import gradio as gr

from document_processing import DocumentError, prepare_document
from parser_client import (
    ParserClient,
    ParserServiceError,
    api_key_name_for_environment,
)
from rate_limit import RateLimiter
from ui_helpers import (
    EMPTY_TABLES,
    EMPTY_TEXT,
    configuration_required_html,
    download_html,
    file_strip_html,
    header_html,
    output_progress_html,
    output_values,
    page_context_html,
    page_limit_html,
    privacy_html,
    progress_html,
    settings_html,
    status_html,
    upload_intro_html,
    viewer_html,
    zoom_label_html,
)

logging.basicConfig(
    level=os.environ.get("LOG_LEVEL", "INFO"),
    format="%(asctime)s %(levelname)s %(message)s",
)
LOGGER = logging.getLogger("cohere-parse")
RATE_LIMITER = RateLimiter(limit=5, window_seconds=10 * 60)
ASSETS_DIR = Path(__file__).parent / "assets"
DEFAULT_ZOOM = 50

gr.set_static_paths(paths=[ASSETS_DIR])


def _request_key(request: gr.Request | None) -> str:
    if request:
        forwarded = request.headers.get("x-forwarded-for", "")
        if forwarded:
            return forwarded.split(",", 1)[0].strip()
        real_ip = request.headers.get("x-real-ip", "").strip()
        if real_ip:
            return real_ip
        if request.client and request.client.host:
            return request.client.host
    return "unknown"


def _view_values(
    state: dict,
    show_boxes: bool,
    render_tables: bool,
    zoom: int,
) -> tuple:
    text, tables, raw = output_values(state, render_tables=render_tables)
    page_warning = page_limit_html(state)
    return (
        viewer_html(state, show_boxes=show_boxes, zoom=zoom),
        file_strip_html(state),
        gr.HTML(value=page_warning, visible=bool(page_warning)),
        gr.Textbox(value=str(state["current_page"] + 1)),
        page_context_html(state),
        zoom_label_html(zoom),
        text,
        tables,
        raw,
        download_html(state),
    )


def prepare_upload(
    file_path: str | None,
    progress=gr.Progress(),
) -> tuple:
    if not file_path:
        raise gr.Error("Choose a PDF or image to continue.")

    progress(0, desc="Preparing document")
    try:
        document = prepare_document(
            file_path,
            delete_source=True,
            enforce_temp_location=True,
            on_progress=lambda current, total: progress(
                (current, total),
                desc=f"Rendering page {current} of {total}",
            ),
        )
    except DocumentError as exc:
        raise gr.Error(str(exc)) from exc

    state = document.to_session()
    view = _view_values(state, False, True, DEFAULT_ZOOM)
    progress(1, desc="Document ready")
    LOGGER.info(
        "document prepared kind=%s bytes=%s total_pages=%s processed_pages=%s",
        state["kind"],
        state["size"],
        state["total_pages"],
        len(state["pages"]),
    )
    return (
        state,
        gr.Column(visible=False),
        gr.Column(visible=True),
        *view,
        settings_html(state),
        gr.HTML(value="", visible=False),
        gr.Button(interactive=True),
        gr.HTML(value="", visible=False),
    )


def begin_upload(file_path: str | None):
    if not file_path:
        return gr.HTML(value="", visible=False)
    return gr.HTML(
        value=progress_html("Preparing document", upload=True),
        visible=True,
    )


def clear_upload_progress():
    return gr.HTML(value="", visible=False)


def begin_run(state: dict | None) -> tuple:
    if not state or not state.get("pages"):
        raise gr.Error("Upload a document first.")

    next_page = next(
        (
            index + 1
            for index, result in enumerate(state["results"])
            if result is None
        ),
        state["current_page"] + 1,
    )
    return (
        gr.HTML(value="", visible=False),
        gr.Button(interactive=False),
        output_progress_html(next_page, len(state["pages"])),
        gr.Tabs(visible=True),
    )


def clear_run_progress() -> tuple:
    return (
        gr.HTML(value="", visible=False),
        gr.Button(interactive=True),
        EMPTY_TEXT,
        gr.Tabs(visible=False),
    )


def run_document(
    state: dict | None,
    show_boxes: bool,
    render_tables: bool,
    zoom: int,
    request: gr.Request,
    progress=gr.Progress(),
):
    if not state or not state.get("pages"):
        raise gr.Error("Upload a document first.")

    allowed, retry_after = RATE_LIMITER.take(_request_key(request))
    if not allowed:
        raise gr.Error(
            f"Too many documents have been submitted. Try again in {retry_after} seconds."
        )

    try:
        client = ParserClient()
    except ParserServiceError as exc:
        raise gr.Error(str(exc)) from exc

    page_count = len(state["pages"])
    parsed_any = False
    for index, page in enumerate(state["pages"]):
        if state["results"][index] is not None:
            continue

        parsed_any = True
        state["current_page"] = index
        progress(
            (index, page_count),
            desc=f"Parsing page {index + 1} of {page_count}",
        )
        try:
            result = client.parse_page(page["png"])
        except ParserServiceError as exc:
            LOGGER.warning(
                "parse failed page=%s completed=%s error=%s",
                index + 1,
                sum(item is not None for item in state["results"]),
                exc,
            )
            view = _view_values(state, show_boxes, render_tables, zoom)
            yield (
                state,
                *view,
                gr.HTML(
                    value=status_html(state, error=str(exc)),
                    visible=True,
                ),
                gr.Button(interactive=True),
            )
            return

        state["results"][index] = result
        LOGGER.info(
            "page parsed page=%s elapsed_ms=%s chars=%s boxes=%s request_id=%s",
            index + 1,
            result["elapsed_ms"],
            result["output_chars"],
            len(result["boxes"]),
            result["request_id"],
        )
        view = _view_values(state, show_boxes, render_tables, zoom)
        yield (
            state,
            *view,
            gr.HTML(
                value=status_html(
                    state,
                    running=index + 1 < page_count,
                    running_page=index + 2,
                ),
                visible=True,
            ),
            gr.Button(interactive=False),
        )

    if not parsed_any and all(state["results"]):
        view = _view_values(state, show_boxes, render_tables, zoom)
        yield (
            state,
            *view,
            gr.HTML(value=status_html(state), visible=True),
            gr.Button(interactive=False),
        )
    progress(1, desc="Document parsed")


def select_page(
    state: dict | None,
    page_value: str,
    show_boxes: bool,
    render_tables: bool,
    zoom: int,
) -> tuple:
    if not state:
        return gr.skip()
    try:
        requested_page = int(page_value)
    except (TypeError, ValueError):
        requested_page = state["current_page"] + 1
    index = max(0, min(len(state["pages"]) - 1, requested_page - 1))
    state["current_page"] = index
    view = _view_values(state, show_boxes, render_tables, zoom)
    return state, *view


def move_page(
    state: dict | None,
    delta: int,
    show_boxes: bool,
    render_tables: bool,
    zoom: int,
) -> tuple:
    if not state:
        return gr.skip()
    index = max(
        0,
        min(len(state["pages"]) - 1, state["current_page"] + delta),
    )
    state["current_page"] = index
    view = _view_values(state, show_boxes, render_tables, zoom)
    return state, *view


def refresh_viewer(state: dict | None, show_boxes: bool, zoom: int):
    if not state:
        return gr.skip()
    return viewer_html(state, show_boxes=show_boxes, zoom=zoom)


def refresh_tables(state: dict | None, render_tables: bool):
    if not state:
        return EMPTY_TABLES
    return output_values(state, render_tables=render_tables)[1]


def change_zoom(
    state: dict | None,
    zoom: int,
    delta: int,
    show_boxes: bool,
) -> tuple:
    next_zoom = max(10, min(100, zoom + delta))
    if not state:
        return next_zoom, gr.skip(), zoom_label_html(next_zoom)
    return (
        next_zoom,
        viewer_html(state, show_boxes=show_boxes, zoom=next_zoom),
        zoom_label_html(next_zoom),
    )


def reset_workspace() -> tuple:
    return (
        None,
        None,
        gr.Column(visible=True),
        gr.Column(visible=False),
        "",
        "",
        gr.HTML(value="", visible=False),
        EMPTY_TEXT,
        EMPTY_TABLES,
        "",
        "",
        gr.HTML(value="", visible=False),
        gr.Textbox(value=""),
        "",
        zoom_label_html(DEFAULT_ZOOM),
        gr.Button(interactive=True),
        False,
        True,
        DEFAULT_ZOOM,
        gr.Tabs(visible=False),
    )


def build_demo() -> gr.Blocks:
    with gr.Blocks(
        title="Cohere Parse",
        fill_height=True,
        fill_width=True,
        delete_cache=(60, 600),
        analytics_enabled=False,
    ) as demo:
        session = gr.State(value=None, time_to_live=15 * 60)
        zoom_level = gr.State(value=DEFAULT_ZOOM)

        gr.HTML(header_html(), elem_id="header-shell")
        try:
            required_api_key = api_key_name_for_environment()
        except ParserServiceError:
            required_api_key = None
        if not required_api_key or not os.environ.get(required_api_key, "").strip():
            gr.HTML(configuration_required_html(required_api_key))

        with gr.Column(elem_id="upload-screen") as upload_panel:
            with gr.Column(elem_id="upload-content"):
                gr.HTML(upload_intro_html())
                upload = gr.File(
                    label="Upload file",
                    file_types=[".pdf", ".png", ".jpg", ".jpeg", ".webp"],
                    file_count="single",
                    type="filepath",
                    elem_id="document-upload",
                )
                upload_progress = gr.HTML(
                    value="",
                    visible=False,
                    elem_id="upload-progress",
                )
            gr.HTML(privacy_html(), elem_id="privacy-note-shell")

        with gr.Column(visible=False, elem_id="workspace") as workspace:
            with gr.Row(equal_height=True, elem_id="workspace-grid"):
                with gr.Column(scale=3, min_width=420, elem_id="document-pane"):
                    file_strip = gr.HTML(elem_id="file-strip-shell")
                    page_warning = gr.HTML(visible=False, elem_id="page-warning")
                    with gr.Row(elem_id="viewer-toolbar"):
                        previous = gr.Button(
                            "←",
                            size="sm",
                            elem_classes=[
                                "icon-button",
                                "page-arrow",
                                "previous-arrow",
                            ],
                        )
                        page_select = gr.Textbox(
                            value="",
                            show_label=False,
                            container=False,
                            max_lines=1,
                            elem_classes="page-select",
                        )
                        page_total = gr.HTML("", elem_classes="page-total")
                        next_page = gr.Button(
                            "→",
                            size="sm",
                            elem_classes=[
                                "icon-button",
                                "page-arrow",
                                "next-arrow",
                            ],
                        )
                        gr.HTML(
                            '<span class="toolbar-spacer"></span>',
                            elem_classes="toolbar-spacer-shell",
                        )
                        zoom_out = gr.Button(
                            "Zoom out",
                            size="sm",
                            elem_classes=["icon-button", "zoom-button", "zoom-out"],
                        )
                        zoom_value = gr.HTML(
                            zoom_label_html(DEFAULT_ZOOM),
                            elem_classes="zoom-value-shell",
                        )
                        zoom_in = gr.Button(
                            "Zoom in",
                            size="sm",
                            elem_classes=["icon-button", "zoom-button", "zoom-in"],
                        )
                    viewer = gr.HTML(elem_classes="document-viewer")

                with gr.Column(scale=2, min_width=380, elem_id="result-pane"):
                    settings = gr.HTML()
                    run_status = gr.HTML(visible=False, elem_id="run-status-shell")
                    with gr.Column(elem_id="parse-controls"):
                        show_boxes = gr.Checkbox(
                            value=False,
                            label="Bounding boxes",
                        )
                        render_tables = gr.Checkbox(
                            value=True,
                            label="Rendered tables",
                        )
                    with gr.Row(elem_id="workspace-actions"):
                        start_over = gr.Button(
                            "Start over",
                            variant="secondary",
                            elem_classes="start-over-button",
                        )
                        download = gr.HTML(elem_classes="download-container")
                        run_button = gr.Button(
                            "Run",
                            variant="primary",
                            elem_classes="run-button",
                        )
                    with gr.Tabs(visible=False, elem_id="output-tabs") as output_tabs:
                        with gr.Tab("Text-only output"):
                            text_output = gr.HTML(
                                value=EMPTY_TEXT,
                                elem_classes="text-output",
                            )
                        with gr.Tab("Tables"):
                            table_output = gr.HTML(
                                value=EMPTY_TABLES,
                                elem_classes="visual-output",
                            )
                        with gr.Tab("Raw"):
                            raw_output = gr.Code(
                                value="",
                                language="markdown",
                                interactive=False,
                                show_label=False,
                                elem_classes="raw-output",
                            )

        prepare_outputs = [
            session,
            upload_panel,
            workspace,
            viewer,
            file_strip,
            page_warning,
            page_select,
            page_total,
            zoom_value,
            text_output,
            table_output,
            raw_output,
            download,
            settings,
            run_status,
            run_button,
            upload_progress,
        ]
        upload_start = upload.upload(
            fn=begin_upload,
            inputs=upload,
            outputs=upload_progress,
            api_visibility="private",
            queue=False,
            show_progress="hidden",
        )
        prepare_event = upload_start.then(
            fn=prepare_upload,
            inputs=upload,
            outputs=prepare_outputs,
            api_visibility="private",
            concurrency_limit=2,
            show_progress="hidden",
        )
        prepare_event.failure(
            fn=clear_upload_progress,
            inputs=None,
            outputs=upload_progress,
            api_visibility="private",
            queue=False,
            show_progress="hidden",
        )

        view_outputs = [
            session,
            viewer,
            file_strip,
            page_warning,
            page_select,
            page_total,
            zoom_value,
            text_output,
            table_output,
            raw_output,
            download,
        ]
        run_outputs = [
            session,
            viewer,
            file_strip,
            page_warning,
            page_select,
            page_total,
            zoom_value,
            text_output,
            table_output,
            raw_output,
            download,
            run_status,
            run_button,
        ]
        run_start = run_button.click(
            fn=begin_run,
            inputs=session,
            outputs=[run_status, run_button, text_output, output_tabs],
            api_visibility="private",
            queue=False,
            show_progress="hidden",
        )
        run_event = run_start.then(
            fn=run_document,
            inputs=[session, show_boxes, render_tables, zoom_level],
            outputs=run_outputs,
            api_visibility="private",
            concurrency_limit=2,
            concurrency_id="parser-api",
            show_progress="hidden",
        )
        run_event.failure(
            fn=clear_run_progress,
            inputs=None,
            outputs=[run_status, run_button, text_output, output_tabs],
            api_visibility="private",
            queue=False,
            show_progress="hidden",
        )

        page_select.submit(
            fn=select_page,
            inputs=[session, page_select, show_boxes, render_tables, zoom_level],
            outputs=view_outputs,
            api_visibility="private",
            queue=False,
        )
        previous.click(
            fn=lambda state, boxes, tables, scale: move_page(
                state, -1, boxes, tables, scale
            ),
            inputs=[session, show_boxes, render_tables, zoom_level],
            outputs=view_outputs,
            api_visibility="private",
            queue=False,
        )
        next_page.click(
            fn=lambda state, boxes, tables, scale: move_page(
                state, 1, boxes, tables, scale
            ),
            inputs=[session, show_boxes, render_tables, zoom_level],
            outputs=view_outputs,
            api_visibility="private",
            queue=False,
        )
        show_boxes.change(
            fn=refresh_viewer,
            inputs=[session, show_boxes, zoom_level],
            outputs=viewer,
            api_visibility="private",
            queue=False,
        )
        zoom_in.click(
            fn=lambda state, scale, boxes: change_zoom(state, scale, 10, boxes),
            inputs=[session, zoom_level, show_boxes],
            outputs=[zoom_level, viewer, zoom_value],
            api_visibility="private",
            queue=False,
        )
        zoom_out.click(
            fn=lambda state, scale, boxes: change_zoom(state, scale, -10, boxes),
            inputs=[session, zoom_level, show_boxes],
            outputs=[zoom_level, viewer, zoom_value],
            api_visibility="private",
            queue=False,
        )
        render_tables.change(
            fn=refresh_tables,
            inputs=[session, render_tables],
            outputs=table_output,
            api_visibility="private",
            queue=False,
        )

        start_over.click(
            fn=reset_workspace,
            inputs=None,
            outputs=[
                session,
                upload,
                upload_panel,
                workspace,
                viewer,
                file_strip,
                page_warning,
                text_output,
                table_output,
                raw_output,
                download,
                run_status,
                page_select,
                page_total,
                zoom_value,
                run_button,
                show_boxes,
                render_tables,
                zoom_level,
                output_tabs,
            ],
            cancels=[run_event],
            api_visibility="private",
            queue=False,
        )

    return demo


demo = build_demo()
demo.queue(
    max_size=20,
    default_concurrency_limit=2,
    api_open=False,
)

if __name__ == "__main__":
    demo.launch(
        css_paths="style.css",
        max_file_size="20mb",
        enable_monitoring=False,
        strict_cors=True,
        blocked_paths=[
            str(Path(__file__).parent / ".env"),
            str(Path(__file__).parent / ".git"),
            str(Path(__file__).parent / "__pycache__"),
        ],
        footer_links=[],
        ssr_mode=False,
    )