import base64
import io
import json
import os
import re
import tempfile
import zipfile
from datetime import datetime
from pathlib import Path

import gradio as gr
import requests
from PIL import Image
import pdf2image

API_URL = "https://t707h6d9q6oftbx3.aistudio-app.com/layout-parsing"
TOKEN = os.getenv("API_TOKEN")


CSS = """
:root {
    --sand-color: #FAF9F6;    
    --white: #ffffff;          
    --shadow: 0 4px 6px rgba(0, 0, 0, 0.1); 
    --text-color: #FAF9F6;     
    --black:#000000;
    --link-hover: #2b6cb0;
    --content-width: 1200px;
}

body {
    display: flex;
    justify-content: center;
    background-color: var(--sand-color);
    color: var(--text-color);
    font-family: Arial, sans-serif;
}

.gradio-container {
    max-width: var(--content-width) !important;
    width: 100% !important;
    margin: 20px auto;
    padding: 20px;
    background-color: var(--white);
}

#component-0, 
#tabs, 
#settings { 
    background-color: var(--white) !important;
    padding: 15px;
}

.upload-section {
    width: 100%;
    margin: 0 auto 30px;
    padding: 20px;
    background-color: var(--sand-color) !important;
    border-radius: 8px;
    box-shadow: var(--shadow);
}

.center-content {
    display: flex;
    flex-direction: column;
    align-items: center;
    text-align: center;
    margin-bottom: 20px;
}

.header {
    margin-bottom: 30px;
    width: 100%;
}

.logo-container {
    width: 100%;
    margin-bottom: 20px;
}

.logo-img {
    width: 100%;
    max-width: var(--content-width);
    margin: 0 auto;
    display: block;
}

.nav-bar {
    display: flex;
    justify-content: center;
    background-color: var(--white);
    padding: 15px 0;
    box-shadow: var(--shadow);
    margin-bottom: 20px;
}

.nav-links {
    display: flex;
    gap: 30px;
    width: 100%;
    justify-content: center;
}

.nav-link {
    color: var(--black);
    text-decoration: none;
    font-weight: bold;
    font-size: 24px;
    transition: color 0.2s;
}

.nav-link:hover {
    color: var(--link-hover);
    text-decoration: none;
}

.result-container {
    display: flex;
    gap: 20px;
    margin-bottom: 30px;
    width: 100%;
}

.pdf-preview {
    flex: 1;
    min-width: 0;
}

.markdown-result {
    flex: 1;
    min-width: 0;
}

.gallery-container {
    width: 100% !important;
}

.gallery-item {
    width: 100% !important;
    height: auto !important;
    aspect-ratio: auto !important;
}

button {
    background-color: var(--text-color) !important;
    color: var(--black) !important;
    border: none !important;
    border-radius: 4px;
    padding: 8px 16px;
}
button:hover {
    opacity: 0.8 !important;
}

.radio-group {
    margin-bottom: 15px !important;
}

.file-download {
    margin-top: 15px !important;
}
.loader {
    border: 5px solid #f3f3f3;
    border-top: 5px solid #3498db;
    border-radius: 50%;
    width: 50px;
    height: 50px;
    animation: spin 1s linear infinite;
    margin: 20px auto;
}

@keyframes spin {
    0% { transform: rotate(0deg); }
    100% { transform: rotate(360deg); }
}

.loader-container {
    text-align: center;
    margin: 20px 0;
}
"""

def clean_markdown_text(text):
    if not text:
        return ""
    text = re.sub(r'<[^>]+>', '', text)
    text = re.sub(r'\n{3,}', '\n\n', text)
    return text.strip()


def pdf_to_images(pdf_path):
    try:
        images = pdf2image.convert_from_path(pdf_path)
        return [image for image in images]
    except:
        return None


def process_file(file_path, file_type):
    try:
        with open(file_path, "rb") as f:
            file_bytes = f.read()

        file_data = base64.b64encode(file_bytes).decode("ascii")
        headers = {
            "Authorization": f"token {TOKEN}",
            "Content-Type": "application/json"
        }

        response = requests.post(
            API_URL,
            json={"file": file_data, "fileType": 0 if file_type == "pdf" else 1},
            headers=headers,
            timeout=60
        )
        response.raise_for_status()

        result = response.json()
        layout_results = result.get("result", {}).get("layoutParsingResults", [])

        markdown_contents = []
        clean_markdown_contents = []
        for res in layout_results:
            markdown = res.get("markdown", {})
            if isinstance(markdown, str):
                original = markdown
            elif isinstance(markdown, dict):
                original = markdown.get("text", "")

            markdown_contents.append(original)
            clean_markdown_contents.append(clean_markdown_text(original))

        if file_type == "pdf":
            images = pdf_to_images(file_path)
        else:
            images = [Image.open(file_path)]

        return {
            "original_file": file_path,
            "markdown_contents": markdown_contents,
            "clean_markdown_contents": clean_markdown_contents,
            "pdf_images": images,
            "api_response": result
        }

    except Exception as e:
        raise gr.Error(f"Error processing file: {str(e)}")


def create_zip_file(results):
    try:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        zip_filename = f"analysis_results_{timestamp}.zip"

        temp_dir = tempfile.mkdtemp()
        zip_path = os.path.join(temp_dir, zip_filename)

        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            original_path = results.get("original_file", "")
            if original_path and Path(original_path).exists():
                zipf.write(original_path, f"original/{Path(original_path).name}")

            markdowns = results.get("markdown_contents", [])
            for i, md_content in enumerate(markdowns):
                if md_content:
                    zipf.writestr(f"markdown/original/markdown_{i + 1}.md", md_content)

            api_response = results.get("api_response", {})
            zipf.writestr("api_response.json", json.dumps(api_response, indent=2, ensure_ascii=False))

        return zip_path

    except Exception as e:
        raise gr.Error(f"Error creating ZIP file: {str(e)}")


def export_markdown(results):
    try:
        markdowns = results.get("markdown_contents", [])
        if not markdowns:
            raise gr.Error("No markdown content to export")

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"original_markdown_{timestamp}.md"
        content = "\n\n".join(markdowns)

        temp_dir = tempfile.mkdtemp()
        file_path = os.path.join(temp_dir, filename)

        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(content)

        return file_path

    except Exception as e:
        raise gr.Error(f"Error exporting markdown: {str(e)}")

with gr.Blocks(css=CSS, title="Document Analysis System") as demo:
    results_state = gr.State()

    # 添加logo图片
    with gr.Column(elem_classes=["logo-container"]):
        gr.Image("英文.png", elem_classes=["logo-img"], show_label=False)

    # 添加导航栏链接
    with gr.Row(elem_classes=["nav-bar"]):
        gr.HTML("""
        <div class="nav-links">
            <a href="https://github.com/PaddlePaddle/PaddleOCR" class="nav-link" target="_blank">GitHub</a>
            <a href="https://paddleocr.ai" class="nav-link" target="_blank">Homepage</a>
        </div>
        """)

    with gr.Column(elem_classes=["upload-section"]):
        file_type = gr.Radio(
            ["pdf", "image"],
            label="File type",
            value="pdf",
            interactive=True
        )
        file_input = gr.File(
            label="Upload document",
            file_types=[".pdf", ".jpg", ".jpeg", ".png"],
            type="filepath"
        )
        process_btn = gr.Button("Analyze document", variant="primary")

        loading_spinner = gr.Column(
            visible=False,
            elem_classes=["loader-container"]
        )
        with loading_spinner:
            gr.HTML("""
            <div class="loader"></div>
            <p>Wait...</p>
            """)

    with gr.Row(elem_classes=["result-container"]):
        with gr.Column(elem_classes=["pdf-preview"]):
            gr.Markdown("### Original document preview")
            pdf_display = gr.Gallery(
                label="PDF page",
                show_label=False,
                elem_classes=["gallery-container"]
            )

        with gr.Column(elem_classes=["markdown-result"]):
            with gr.Row(elem_classes=["radio-group"]):
                display_mode = gr.Radio(
                    ["Original Markdown", "Cleaned Text"],
                    label="Display Mode",
                    value="Original Markdown",
                    interactive=True
                )
            markdown_display = gr.HTML(label="Analysis Results")
    with gr.Column(elem_classes=["download-section"]):
        gr.Markdown("### Result Export")
        with gr.Row():
            download_md_btn = gr.Button("Download Original Markdown", variant="secondary")
            download_all_btn = gr.Button("Download Complete Analysis Results (ZIP)", variant="primary")
        download_file = gr.File(visible=False, label="Download file", elem_classes=["file-download"])

    def toggle_spinner():
        return gr.update(visible=True)


    def hide_spinner():
        return gr.update(visible=False)


    process_btn.click(
        toggle_spinner,
        outputs=[loading_spinner]
    ).then(
        process_file,
        inputs=[file_input, file_type],
        outputs=[results_state]
    ).then(
        hide_spinner,
        outputs=[loading_spinner]
    ).success(
        lambda res: res["pdf_images"] if res and res.get("pdf_images") else [],
        inputs=[results_state],
        outputs=[pdf_display]
    ).success(
        lambda res: res["markdown_contents"][0] if res and res.get("markdown_contents") else "",
        inputs=[results_state],
        outputs=[markdown_display]
    )

    display_mode.change(
        lambda mode, res: (
            res["markdown_contents"][0] if mode == "原始Markdown"
            else res["clean_markdown_contents"][0]
        ) if res else "",
        inputs=[display_mode, results_state],
        outputs=[markdown_display]
    )

    download_md_btn.click(
        export_markdown,
        inputs=[results_state],
        outputs=[download_file]
    ).then(
        lambda x: gr.update(visible=True),
        inputs=[download_file],
        outputs=[download_file]
    )

    download_all_btn.click(
        create_zip_file,
        inputs=[results_state],
        outputs=[download_file]
    ).then(
        lambda x: gr.update(visible=True),
        inputs=[download_file],
        outputs=[download_file]
    )

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