import json
import tempfile
from pathlib import Path

import gradio as gr
import spaces

from document_processor import DocumentProcessor
from normalizer import DocumentNormalizer

processor = DocumentProcessor()
normalizer = DocumentNormalizer()

OUTPUT_DIR = Path(tempfile.gettempdir()) / "ai-tutor-output"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)


@spaces.GPU(duration=120)
def process_document(file):

    if file is None:
        raise gr.Error("Please upload a document.")

    file_path = file

    if hasattr(file, "name"):
        file_path = file.name

    path = Path(file_path)

    try:
        result = processor.process(str(path))

        normalized = normalizer.normalize(
            docling_document=result["document"],
            filename=path.name,
            raw_json=result["raw_json"],
        )

        output = normalized.model_dump()

        base_name = path.stem
        json_path = OUTPUT_DIR / f"{base_name}_normalized.json"
        markdown_path = OUTPUT_DIR / f"{base_name}_markdown.md"
        stats_path = OUTPUT_DIR / f"{base_name}_statistics.json"

        json_path.write_text(
            json.dumps(output, indent=2, ensure_ascii=False),
            encoding="utf-8",
        )
        markdown_path.write_text(result["markdown"], encoding="utf-8")
        stats_path.write_text(
            json.dumps(normalized.statistics, indent=2, ensure_ascii=False),
            encoding="utf-8",
        )

        download_files = [
            str(json_path),
            str(markdown_path),
            str(stats_path),
        ]

        return (
            json.dumps(
                output,
                indent=2,
                ensure_ascii=False,
            ),
            result["markdown"],
            json.dumps(
                normalized.statistics,
                indent=2,
                ensure_ascii=False,
            ),
            download_files,
        )

    except Exception as e:
        raise gr.Error(
            f"Document processing failed: {type(e).__name__}: {e}"
        )


with gr.Blocks(
    title="AI Tutor Document Processor"
) as demo:

    gr.Markdown(
        """
        # 📚 AI Tutor Document Processor

        Convert educational documents into a structured
        AI-Tutor document model.

        **Supported:** PDF, DOCX, PPTX, HTML, Markdown, TXT
        """
    )

    with gr.Row():

        file_input = gr.File(
            label="Upload Document",
            file_types=[
                ".pdf",
                ".docx",
                ".pptx",
                ".html",
                ".htm",
                ".md",
                ".txt",
            ],
            type="filepath",
        )

        process_button = gr.Button(
            "Process Document",
            variant="primary",
        )

    gr.Markdown("## Processing Statistics")

    statistics_output = gr.Code(
        label="Statistics",
        language="json",
    )

    gr.Markdown("## Download")

    download_output = gr.File(
        label="Download Processed Files",
        file_count="multiple",
    )

    with gr.Tabs():

        with gr.Tab("AI Tutor JSON"):

            json_output = gr.Code(
                label="Normalized Document",
                language="json",
            )

        with gr.Tab("Markdown"):

            markdown_output = gr.Markdown()

    process_button.click(
        fn=process_document,
        inputs=file_input,
        outputs=[
            json_output,
            markdown_output,
            statistics_output,
            download_output,
        ],
        api_name="process_document",
    )


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