import gradio as gr
import re
from typing import Dict, Any, Optional
import os
from pathlib import Path

# Import dependencies for PDF and DOCX processing
try:
    import docx

    DOCX_AVAILABLE = True
except ImportError:
    DOCX_AVAILABLE = False

try:
    import fitz  # PyMuPDF

    PDF_AVAILABLE = True
except ImportError:
    PDF_AVAILABLE = False


class DocumentToMarkdownConverter:
    def __init__(self):
        self.elements = []

    def extract_from_docx(self, docx_path: str) -> str:
        """Extract content from DOCX and convert to Markdown"""
        if not DOCX_AVAILABLE:
            raise ImportError("python-docx not installed. Run: pip install python-docx")

        doc = docx.Document(docx_path)
        markdown_content = []

        # Process paragraphs
        for paragraph in doc.paragraphs:
            if paragraph.text.strip():
                md_text = self._convert_paragraph_to_markdown(paragraph)
                if md_text:
                    markdown_content.append(md_text)

        # Process tables
        for table in doc.tables:
            md_table = self._convert_table_to_markdown(table)
            if md_table:
                markdown_content.append(md_table)

        return "\n\n".join(markdown_content)

    def extract_from_pdf(self, pdf_path: str) -> str:
        """Extract content from PDF and convert to Markdown"""
        if not PDF_AVAILABLE:
            raise ImportError("PyMuPDF not installed. Run: pip install PyMuPDF")

        doc = fitz.open(pdf_path)
        markdown_content = []

        for page_num in range(len(doc)):
            page = doc.load_page(page_num)

            # Extract text blocks with formatting
            blocks = page.get_text("dict")
            page_markdown = self._convert_pdf_blocks_to_markdown(blocks)

            if page_markdown.strip():
                markdown_content.append(f"## Page {page_num + 1}\n\n{page_markdown}")

        doc.close()
        return "\n\n---\n\n".join(markdown_content)

    def _convert_paragraph_to_markdown(self, paragraph) -> str:
        """Convert DOCX paragraph to Markdown"""
        text = paragraph.text.strip()
        if not text:
            return ""

        style_name = paragraph.style.name if paragraph.style else "Normal"

        # Check if paragraph has bold formatting
        is_bold = any(run.bold for run in paragraph.runs if run.bold)

        # Check font size for heading detection
        font_size = 12
        if paragraph.runs:
            first_run = paragraph.runs[0]
            if first_run.font.size:
                font_size = first_run.font.size.pt

        # Convert based on style and formatting
        if "Title" in style_name or (is_bold and font_size >= 18):
            return f"# {text}"
        elif "Heading 1" in style_name or (is_bold and font_size >= 16):
            return f"# {text}"
        elif "Heading 2" in style_name or (is_bold and font_size >= 14):
            return f"## {text}"
        elif "Heading 3" in style_name or (is_bold and font_size >= 13):
            return f"### {text}"
        elif "Heading 4" in style_name:
            return f"#### {text}"
        elif "Heading 5" in style_name:
            return f"##### {text}"
        elif "Heading 6" in style_name:
            return f"###### {text}"
        elif re.match(r"^[\d\w]\.\s|^[•\-\*]\s|^\d+\)\s", text):
            # List items
            if text.startswith(("1.", "2.", "3.", "4.", "5.", "6.", "7.", "8.", "9.")):
                return f"1. {text[2:].strip()}"
            else:
                return f"- {text[1:].strip() if text[0] in '•-*' else text}"
        else:
            # Regular paragraph
            formatted_text = self._apply_inline_formatting(paragraph)
            return formatted_text

    def _apply_inline_formatting(self, paragraph) -> str:
        """Apply inline formatting (bold, italic) to text"""
        result = ""
        for run in paragraph.runs:
            text = run.text
            if run.bold and run.italic:
                text = f"***{text}***"
            elif run.bold:
                text = f"**{text}**"
            elif run.italic:
                text = f"*{text}*"
            result += text
        return result

    def _convert_table_to_markdown(self, table) -> str:
        """Convert DOCX table to Markdown table"""
        if not table.rows:
            return ""

        markdown_rows = []

        # Process header row
        header_cells = [cell.text.strip() for cell in table.rows[0].cells]
        markdown_rows.append("| " + " | ".join(header_cells) + " |")
        markdown_rows.append("| " + " | ".join(["---"] * len(header_cells)) + " |")

        # Process data rows
        for row in table.rows[1:]:
            cells = [cell.text.strip() for cell in row.cells]
            markdown_rows.append("| " + " | ".join(cells) + " |")

        return "\n".join(markdown_rows)

    def _convert_pdf_blocks_to_markdown(self, blocks_dict) -> str:
        """Convert PDF text blocks to Markdown"""
        markdown_lines = []

        for block in blocks_dict.get("blocks", []):
            if block.get("type") == 0:  # Text block
                for line in block.get("lines", []):
                    line_text = ""
                    for span in line.get("spans", []):
                        text = span.get("text", "").strip()
                        if text:
                            # Check formatting
                            font_size = span.get("size", 12)
                            flags = span.get("flags", 0)

                            # Bold = flags & 16, Italic = flags & 2
                            is_bold = bool(flags & 16)
                            is_italic = bool(flags & 2)

                            # Apply formatting
                            if is_bold and is_italic:
                                text = f"***{text}***"
                            elif is_bold:
                                text = f"**{text}**"
                            elif is_italic:
                                text = f"*{text}*"

                            # Check if it's a heading based on font size
                            if font_size >= 18:
                                text = f"# {text}"
                            elif font_size >= 16:
                                text = f"## {text}"
                            elif font_size >= 14:
                                text = f"### {text}"

                            line_text += text + " "

                    if line_text.strip():
                        markdown_lines.append(line_text.strip())

        return "\n\n".join(markdown_lines)

    def analyze_markdown_structure(self, markdown_text: str) -> Dict[str, Any]:
        """Analyze the structure of extracted Markdown"""
        lines = markdown_text.split("\n")
        structure = {
            "headings": {"h1": 0, "h2": 0, "h3": 0, "h4": 0, "h5": 0, "h6": 0},
            "lists": {"ordered": 0, "unordered": 0},
            "tables": 0,
            "paragraphs": 0,
            "bold_text": 0,
            "italic_text": 0,
            "total_lines": len(lines),
            "word_count": len(markdown_text.split()),
            "character_count": len(markdown_text),
        }

        in_table = False

        for line in lines:
            line = line.strip()
            if not line:
                continue

            # Count headings
            if line.startswith("#"):
                level = len(line) - len(line.lstrip("#"))
                if level <= 6:
                    structure["headings"][f"h{level}"] += 1

            # Count lists
            elif re.match(r"^\d+\.\s", line):
                structure["lists"]["ordered"] += 1
            elif re.match(r"^[\-\*\+]\s", line):
                structure["lists"]["unordered"] += 1

            # Count tables
            elif "|" in line and not in_table:
                structure["tables"] += 1
                in_table = True
            elif "|" not in line:
                in_table = False
                if (
                    line
                    and not line.startswith("#")
                    and not re.match(r"^[\-\*\+\d]", line)
                ):
                    structure["paragraphs"] += 1

            # Count formatting
            structure["bold_text"] += len(re.findall(r"\*\*[^*]+\*\*", line))
            structure["italic_text"] += len(re.findall(r"\*[^*]+\*", line))

        return structure


def extract_document_to_markdown(file_path: str) -> Dict[str, Any]:
    """
    Extract document content and convert to Markdown format

    Args:
        file_path: Path to PDF or DOCX file

    Returns:
        Dictionary containing markdown content and structure analysis
    """

    if not file_path or not os.path.exists(file_path):
        return {"error": "File not found", "markdown": "", "structure": {}}

    converter = DocumentToMarkdownConverter()
    file_extension = Path(file_path).suffix.lower()

    try:
        if file_extension == ".docx":
            if not DOCX_AVAILABLE:
                return {
                    "error": "python-docx not installed. Run: pip install python-docx",
                    "markdown": "",
                    "structure": {},
                }
            markdown_content = converter.extract_from_docx(file_path)

        elif file_extension == ".pdf":
            if not PDF_AVAILABLE:
                return {
                    "error": "PyMuPDF not installed. Run: pip install PyMuPDF",
                    "markdown": "",
                    "structure": {},
                }
            markdown_content = converter.extract_from_pdf(file_path)

        else:
            return {
                "error": f"Unsupported file type: {file_extension}. Only PDF and DOCX files are supported.",
                "markdown": "",
                "structure": {},
            }

        # Analyze markdown structure
        structure = converter.analyze_markdown_structure(markdown_content)

        return {
            "success": True,
            "file_info": {
                "name": Path(file_path).name,
                "type": file_extension.upper()[1:],
                "size_kb": round(os.path.getsize(file_path) / 1024, 2),
            },
            "markdown": markdown_content,
            "structure": structure,
            "preview": markdown_content[:500] + "..."
            if len(markdown_content) > 500
            else markdown_content,
        }

    except Exception as e:
        return {
            "error": f"Error processing file: {str(e)}",
            "markdown": "",
            "structure": {},
        }


# Create Gradio interface
def create_interface():
    with gr.Blocks(title="Document to Markdown Converter") as demo:
        gr.Markdown("# 📄 Document to Markdown Converter")
        gr.Markdown(
            "Upload PDF or DOCX files to extract content and convert to Markdown format"
        )

        missing_deps = []
        if not DOCX_AVAILABLE:
            missing_deps.append("python-docx")
        if not PDF_AVAILABLE:
            missing_deps.append("PyMuPDF")

        if missing_deps:
            gr.Markdown(
                f"⚠️ **Missing dependencies**: Run `pip install {' '.join(missing_deps)}` to enable full support"
            )

        with gr.Row():
            with gr.Column(scale=1):
                file_input = gr.File(
                    label="Upload Document",
                    file_types=[".pdf", ".docx"],
                    type="filepath",
                )
                extract_btn = gr.Button("Extract to Markdown", variant="primary")

                with gr.Accordion("Output Options", open=False):
                    show_structure = gr.Checkbox(
                        label="Show Structure Analysis", value=True
                    )
                    show_preview = gr.Checkbox(label="Show Preview Only", value=False)

            with gr.Column(scale=2):
                with gr.Tabs():
                    with gr.TabItem("Markdown Output"):
                        markdown_output = gr.Textbox(
                            label="Extracted Markdown",
                            lines=20,
                            max_lines=30,
                            show_copy_button=True,
                        )

                    with gr.TabItem("Structure Analysis"):
                        structure_output = gr.JSON(label="Document Structure")

                    with gr.TabItem("File Info"):
                        info_output = gr.JSON(label="File Information")

        def process_document(file_path, show_struct, show_prev):
            if not file_path:
                return "No file uploaded", {}, {}

            result = extract_document_to_markdown(file_path)

            if "error" in result:
                return f"Error: {result['error']}", {}, {}

            markdown_text = result["preview"] if show_prev else result["markdown"]
            structure = result["structure"] if show_struct else {}
            file_info = result["file_info"]

            return markdown_text, structure, file_info

        extract_btn.click(
            fn=process_document,
            inputs=[file_input, show_structure, show_preview],
            outputs=[markdown_output, structure_output, info_output],
        )

    return demo


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