import gradio as gr
from pathlib import Path
import datetime
import re
import os
import shutil
import io
import base64
from collections import defaultdict
from PIL import Image
import json

# Document Generation Libs
from docx import Document
import openpyxl
from pypdf import PdfWriter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, PageBreak, BaseDocTemplate, Frame, PageTemplate
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.pagesizes import letter, A4, landscape
from reportlab.lib.units import inch
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont

# Media Libs
import fitz  # PyMuPDF

# --- Configuration & Setup ---
CWD = Path.cwd()
OUTPUT_DIR = CWD / "generated_outputs"
PREVIEW_DIR = CWD / "previews"
UPLOAD_DIR = CWD / "uploads"
FONT_DIR = CWD

# Create necessary directories
OUTPUT_DIR.mkdir(exist_ok=True)
PREVIEW_DIR.mkdir(exist_ok=True)
UPLOAD_DIR.mkdir(exist_ok=True)

LAYOUTS = {
    "A4 Portrait": {"size": A4},
    "A4 Landscape": {"size": landscape(A4)},
    "Letter Portrait": {"size": letter},
    "Letter Landscape": {"size": landscape(letter)},
}

# --- ✍️ Document Generation Engines ---

def create_pdf(md_content, font_name, emoji_font, pagesize, num_columns):
    """📄 Builds a beautiful PDF from a Markdown story using ReportLab."""
    pdf_buffer = io.BytesIO()
    story = markdown_to_story(md_content, font_name, emoji_font)
    if num_columns > 1:
        doc = BaseDocTemplate(pdf_buffer, pagesize=pagesize, leftMargin=0.5 * inch, rightMargin=0.5 * inch)
        frame_width = (doc.width / num_columns) - (num_columns - 1) * 0.1 * inch
        frames = [Frame(doc.leftMargin + i * (frame_width + 0.2 * inch), doc.bottomMargin, frame_width, doc.height) for i in range(num_columns)]
        doc.addPageTemplates([PageTemplate(id='MultiCol', frames=frames)])
    else:
        doc = SimpleDocTemplate(pdf_buffer, pagesize=pagesize)
    doc.build(story)
    pdf_buffer.seek(0)
    return pdf_buffer

def create_docx(md_content):
    """📝 Crafts a DOCX document, translating Markdown to Word elements."""
    document = Document()
    for line in md_content.split('\n'):
        if line.startswith('# '): document.add_heading(line[2:], level=1)
        elif line.startswith('## '): document.add_heading(line[3:], level=2)
        elif line.strip().startswith(('- ', '* ')): document.add_paragraph(line.strip()[2:], style='List Bullet')
        else:
            p = document.add_paragraph()
            parts = re.split(r'(\*\*.*?\*\*)', line)
            for part in parts:
                if part.startswith('**') and part.endswith('**'): p.add_run(part[2:-2]).bold = True
                else: p.add_run(part)
    return document

def create_xlsx(md_content):
    """📊 Organizes a Markdown outline into columns in an XLSX file."""
    workbook = openpyxl.Workbook(); sheet = workbook.active
    sections = re.split(r'\n# ', '\n' + md_content)
    if sections and sections[0] == '': sections.pop(0)
    column_data = []
    for section in sections:
        lines = section.split('\n'); header = lines[0]
        content = [l.strip() for l in lines[1:] if l.strip()]
        column_data.append({'header': header, 'content': content})
    for c_idx, col in enumerate(column_data, 1):
        sheet.cell(row=1, column=c_idx, value=col['header'])
        for r_idx, line_content in enumerate(col['content'], 2):
            sheet.cell(row=r_idx, column=c_idx, value=line_content)
    return workbook

def markdown_to_story(markdown_text: str, font_name: str, emoji_font: str):
    """📜 Translates Markdown text into a sequence of ReportLab flowables for PDF rendering."""
    styles = getSampleStyleSheet()
    bold_font = f"{font_name}-Bold" if font_name != "Helvetica" else "Helvetica-Bold"
    style_normal = ParagraphStyle('BodyText', fontName=font_name, spaceAfter=6, fontSize=10, leading=14)
    style_h1 = ParagraphStyle('h1', fontName=bold_font, spaceBefore=12, fontSize=24, textColor=colors.HexColor("#1E3A8A"))
    style_h2 = ParagraphStyle('h2', fontName=bold_font, spaceBefore=10, fontSize=18, textColor=colors.HexColor("#374151"))
    style_h3 = ParagraphStyle('h3', fontName=bold_font, spaceBefore=8, fontSize=14, textColor=colors.HexColor("#4B5563"))
    
    story, first_heading = [], True
    for line in markdown_text.split('\n'):
        stripped_line = line.strip()
        if not stripped_line:
            story.append(Spacer(1, 0.1 * inch)); continue

        content, style, extra_args = stripped_line, style_normal, {}
        if stripped_line.startswith("# "):
            if not first_heading: story.append(PageBreak())
            content, style, first_heading = stripped_line.lstrip('# '), style_h1, False
        elif stripped_line.startswith("## "): content, style = stripped_line.lstrip('## '), style_h2
        elif stripped_line.startswith("### "): content, style = stripped_line.lstrip('### '), style_h3
        elif stripped_line.startswith(("- ", "* ")): content, extra_args = stripped_line[2:], {'bulletText': '•'}
        
        formatted_content = re.sub(r'_(.*?)_', r'<i>\1</i>', re.sub(r'\*\*(.*?)\*\*', r'<b>\1</b>', content))
        final_content = apply_emoji_font(formatted_content, emoji_font)
        
        story.append(Paragraph(final_content, style, **extra_args))
    return story


# --- 🔮 Virtual AI Omni-Model Functions ---

def process_text_input(prompt):
    """💬 Simulates an AI response to a text prompt."""
    return f"# Virtual AI Response\n\n**Your Prompt:**\n> {prompt}\n\n**Generated Content:**\n- This is a simulated response for your text input.\n- Here's an emoji: 😊"

def process_image_input(image_path, prompt):
    """🖼️ Simulates an AI description of an image."""
    return f"# Virtual AI Image Analysis: {Path(image_path).name}\n\n**Your Prompt:**\n> {prompt}\n\n**Generated Content:**\n1. Simulated analysis of the uploaded image.\n2. File type appears to be `{Path(image_path).suffix}`."

def process_audio_input(audio_path, prompt):
    """🎤 Simulates AI transcription and summarization of an audio file."""
    return f"# Virtual AI Audio Summary: {Path(audio_path).name}\n\n**Your Prompt:**\n> {prompt}\n\n**Simulated Transcription:**\n> \"This is a test of the emergency broadcast system.\"\n\n**Generated Summary:**\nThe audio is a test broadcast."

def process_pdf_input(pdf_path, prompt, progress):
    """📄 Simulates AI-powered OCR of a PDF document."""
    progress(0.5, desc="Simulating PDF page processing...")
    ocr_text = f"# Virtual AI OCR of: {Path(pdf_path).name}\n\n**Your Prompt:**\n> {prompt}\n\n**Extracted Content (Simulated):**\n- **Page 1:** Simulated text from the first page.\n- **Page 2:** Simulated text from the second page."
    progress(1.0, desc="PDF OCR Simulation Complete!")
    return ocr_text


# --- 🛠️ Helpers & Main API ---

def register_local_fonts():
    """✒️ Scans for local .ttf fonts and registers them for PDF creation."""
    text_font_names, emoji_font_name = [], None
    font_files = list(FONT_DIR.glob("*.ttf"))
    for font_path in font_files:
        try:
            font_name = font_path.stem
            pdfmetrics.registerFont(TTFont(font_name, str(font_path)))
            pdfmetrics.registerFont(TTFont(f"{font_name}-Bold", str(font_path)))
            pdfmetrics.registerFontFamily(font_name, normal=font_name, bold=f"{font_name}-Bold")
            if "notocoloremoji-regular" in font_name.lower():
                emoji_font_name = font_name
            else:
                text_font_names.append(font_name)
        except: pass
    if not text_font_names: text_font_names.append('Helvetica')
    return sorted(text_font_names), emoji_font_name

def apply_emoji_font(text: str, emoji_font_name: str) -> str:
    """😊 Finds emojis and wraps them in special font tags for the PDF."""
    if not emoji_font_name: return text
    emoji_pattern = re.compile(f"([{re.escape(''.join(map(chr, range(0x1f600, 0x1f650))))}"
                               f"{re.escape(''.join(map(chr, range(0x1f300, 0x1f5ff))))}]+)")
    return emoji_pattern.sub(fr'<font name="{emoji_font_name}">\1</font>', text)

def create_pdf_preview(pdf_path: Path):
    """🏞️ Generates a PNG thumbnail for the first page of a PDF."""
    preview_path = PREVIEW_DIR / f"{pdf_path.stem}.png"
    try:
        doc = fitz.open(pdf_path); page = doc.load_page(0); pix = page.get_pixmap(dpi=96)
        pix.save(str(preview_path)); doc.close()
        return preview_path
    except: return None

def build_file_explorer_html(generated_files, pdf_files_for_gallery):
    """🗂️ Constructs the HTML/JS for the file explorer and PDF gallery."""
    
    file_explorer_html = ""
    file_icons = {".pdf": "📄", ".docx": "📝", ".xlsx": "📊"}
    action_buttons = """
        <div class="actions">
            <button class="action-btn" onclick="event.preventDefault(); alert('QuizMe: Feature coming soon!')">🧠 QuizMe</button>
            <button class="action-btn" onclick="event.preventDefault(); alert('Revise: Feature coming soon!')">✍️ Revise</button>
        </div>
    """
    for file_path in generated_files:
        icon = file_icons.get(file_path.suffix, '📎')
        file_explorer_html += f"""
        <div class="file-item">
            <a href="/file={file_path}" class="file-link" download="{file_path.name}">
                <span class="file-icon">{icon}</span>
                <span class="file-name">{file_path.name}</span>
            </a>
            {action_buttons if file_path.suffix == '.pdf' else ''}
        </div>
        """

    gallery_items = []
    for pdf_path in pdf_files_for_gallery:
        preview_path = create_pdf_preview(pdf_path)
        if preview_path:
            with open(preview_path, "rb") as f:
                img_base64 = base64.b64encode(f.read()).decode("utf-8")
                gallery_items.append({"preview_src": f"data:image/png;base64,{img_base64}", "filename": pdf_path.name})
    
    gallery_html = ""
    if gallery_items:
        thumbs_html = "".join([f'<img src="{item["preview_src"]}" class="thumbnail" onclick="selectThumbnail(this, \'{item["preview_src"]}\', \'{item["filename"]}\')">' for item in gallery_items])
        gallery_html = f"""
        <div class="gallery-container">
            <div class="main-view"><img id="main-image" src="{gallery_items[0]['preview_src']}" class="main-image"></div>
            <div class="thumbnail-strip">{thumbs_html}</div>
        </div>
        <p id="main-filename">{gallery_items[0]['filename']}</p>
        """
    
    html = f"""
    <style>
        .tabs {{ display: flex; border-bottom: 2px solid #ccc; }}
        .tab-button {{ padding: 10px 15px; cursor: pointer; background: #f1f1f1; border: none; font-size: 1.2em; }}
        .tab-button.active {{ background: #fff; border-top: 2px solid #6366f1; border-left: 1px solid #ccc; border-right: 1px solid #ccc; position: relative; top: 1px;}}
        .tab-content {{ display: none; padding: 15px; border: 1px solid #ccc; border-top: none; }}
        .tab-content.active {{ display: block; }}
        .file-explorer {{ display: grid; grid-template-columns: repeat(auto-fill, minmax(250px, 1fr)); gap: 15px; }}
        .file-item {{ background: #f9f9f9; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1); }}
        .file-link {{ display: flex; align-items: center; padding: 15px; text-decoration: none; color: #333; }}
        .file-icon {{ font-size: 3em; margin-right: 15px; }}
        .file-name {{ font-weight: bold; word-break: break-all; }}
        .actions {{ display: flex; justify-content: space-around; padding: 5px 10px; border-top: 1px solid #eee; }}
        .action-btn {{ background: none; border: none; cursor: pointer; font-size: 1.1em; }}
        .gallery-container {{ display: flex; height: 500px; }}
        .main-view {{ flex: 4; }} .main-image {{ width: 100%; height: 100%; object-fit: contain; }}
        .thumbnail-strip {{ flex: 1; overflow-y: auto; }} .thumbnail {{ width: 100%; margin-bottom: 5px; cursor: pointer; border: 3px solid transparent; }}
        .thumbnail.active {{ border-color: #6366f1; }}
        #main-filename {{ text-align: center; font-weight: bold; margin-top: 10px; }}
    </style>
    <div class="tabs">
        <button class="tab-button active" onclick="openTab(event, 'explorer')">🗂️ FileExplorer</button>
        {'<button class="tab-button" onclick="openTab(event, \'gallery\')">🖼️ GlimpsePDFs</button>' if gallery_items else ''}
    </div>
    <div id="explorer" class="tab-content active"><div class="file-explorer">{file_explorer_html}</div></div>
    <div id="gallery" class="tab-content">{gallery_html}</div>
    <script>
        function openTab(evt, tabName) {{
            let i, tabcontent = document.getElementsByClassName("tab-content"), tablinks = document.getElementsByClassName("tab-button");
            for (i = 0; i < tabcontent.length; i++) tabcontent[i].style.display = "none";
            for (i = 0; i < tablinks.length; i++) tablinks[i].className = tablinks[i].className.replace(" active", "");
            document.getElementById(tabName).style.display = "block"; evt.currentTarget.className += " active";
        }}
        const mainImage = document.getElementById('main-image'), mainFilename = document.getElementById('main-filename'), thumbnails = document.querySelectorAll('.thumbnail');
        if (thumbnails.length > 0) thumbnails[0].classList.add('active');
        function selectThumbnail(thumb, src, name) {{
            mainImage.src = src; mainFilename.textContent = name;
            thumbnails.forEach(t => t.classList.remove('active')); thumb.classList.add('active');
        }};
    </script>
    """
    return html

def generate_outputs_api(omni_file, omni_prompt, output_formats, layouts, fonts, num_columns, page_w_mult, page_h_mult, progress=gr.Progress(track_tqdm=True)):
    if not omni_prompt and not omni_file: raise gr.Error("Please provide a prompt or upload at least one file.")
    if not output_formats: raise gr.Error("Please select at least one output format.")
    
    shutil.rmtree(OUTPUT_DIR, ignore_errors=True); shutil.rmtree(PREVIEW_DIR, ignore_errors=True)
    OUTPUT_DIR.mkdir(); PREVIEW_DIR.mkdir()

    md_content = ""
    if omni_file:
        temp_path = UPLOAD_DIR / Path(omni_file.name).name
        shutil.copyfile(omni_file.name, temp_path)
        file_ext = temp_path.suffix.lower()

        if file_ext == '.md': md_content = temp_path.read_text(encoding='utf-8')
        elif file_ext == '.pdf': md_content = process_pdf_input(temp_path, omni_prompt or "Extract text", progress)
        elif file_ext in ['.png', '.jpg', '.jpeg']: md_content = process_image_input(temp_path, omni_prompt or "Describe image")
        elif file_ext in ['.wav', '.mp3']: md_content = process_audio_input(temp_path, omni_prompt or "Summarize transcription")
    elif omni_prompt:
        md_content = process_text_input(omni_prompt)
    
    if not md_content: raise gr.Error("Failed to generate source content.")
    
    generated_files, pdf_files_for_gallery = [], []
    for format_choice in progress.tqdm(output_formats, desc="Generating Formats"):
        time_str = datetime.datetime.now().strftime('%m-%d-%a_%I%M%p').upper()
        if format_choice == "PDF":
            for layout_name in layouts:
                for font_name in fonts:
                    pagesize = LAYOUTS[layout_name]["size"]
                    final_pagesize = (pagesize[0] * page_w_mult, pagesize[1] * page_h_mult)
                    pdf_buffer = create_pdf(md_content, font_name, EMOJI_FONT_NAME, final_pagesize, num_columns)
                    filename = f"Document_{time_str}.pdf"; output_path = OUTPUT_DIR / filename
                    with open(output_path, "wb") as f: f.write(pdf_buffer.getvalue())
                    generated_files.append(output_path); pdf_files_for_gallery.append(output_path)
        elif format_choice == "DOCX":
            doc = create_docx(md_content); filename = f"Document_{time_str}.docx"
            output_path = OUTPUT_DIR / filename; doc.save(output_path); generated_files.append(output_path)
        elif format_choice == "XLSX":
            book = create_xlsx(md_content); filename = f"Outline_{time_str}.xlsx"
            output_path = OUTPUT_DIR / filename; book.save(output_path); generated_files.append(output_path)
            
    final_html_output = build_file_explorer_html(generated_files, pdf_files_for_gallery)
    
    return md_content, final_html_output

# --- 🎨 Gradio UI Definition ---
AVAILABLE_FONTS, EMOJI_FONT_NAME = register_local_fonts()

with gr.Blocks(theme=gr.themes.Soft(), title="Omni-Model Document Generator") as demo:
    gr.Markdown("# 🧠 Omni-Model Document Generator")
    gr.Markdown("Provide a prompt or upload a file (MD, PDF, Image, Audio). A virtual AI will process it, and you can generate documents from the result.")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Input & Settings")
            omni_prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Ask a question or provide instructions...")
            omni_file = gr.File(label="Upload File (Optional)", file_types=["image", ".wav", ".mp3", ".md", ".pdf"])
            
            output_formats = gr.CheckboxGroup(choices=["PDF", "DOCX", "XLSX"], label="Select Output Formats", value=["PDF"])
            
            with gr.Accordion("PDF Customization", open=False):
                num_columns_slider = gr.Slider(label="Text Columns", minimum=1, maximum=4, step=1, value=1)
                page_w_mult_slider = gr.Slider(label="Page Width Multiplier", minimum=1, maximum=5, step=1, value=1)
                page_h_mult_slider = gr.Slider(label="Page Height Multiplier", minimum=1, maximum=2, step=1, value=1)
                selected_layouts = gr.CheckboxGroup(choices=list(LAYOUTS.keys()), label="Base Page Layout", value=["A4 Portrait"])
                selected_fonts = gr.CheckboxGroup(choices=AVAILABLE_FONTS, label="Text Font", value=[AVAILABLE_FONTS[0]] if AVAILABLE_FONTS else [])
            
            generate_btn = gr.Button("🚀 Generate Documents", variant="primary")
        
        with gr.Column(scale=2):
            gr.Markdown("### 🤖 AI Response (Source for Documents)")
            ai_response_output = gr.Markdown(label="AI Generated Content")
            gr.Markdown("### 🗂️ Generated Files")
            file_explorer_output = gr.HTML(label="File Explorer & Gallery")

    generate_btn.click(fn=generate_outputs_api, 
                       inputs=[omni_file, omni_prompt, output_formats, selected_layouts, selected_fonts, num_columns_slider, page_w_mult_slider, page_h_mult_slider], 
                       outputs=[ai_response_output, file_explorer_output])

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