import io
import re
import os
import glob
import asyncio
import hashlib
import unicodedata
import streamlit as st
from PIL import Image
import fitz
import edge_tts
from reportlab.lib.pagesizes import A4
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib import colors
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.pdfgen import canvas
from datetime import datetime
import pytz
from pypdf import PdfReader, PdfWriter
from pypdf.annotations import Link
from pypdf.generic import Fit

st.set_page_config(layout="wide", initial_sidebar_state="expanded")

# Existing functions (unchanged)
def get_timestamp_prefix():
    central = pytz.timezone("US/Central")
    now = datetime.now(central)
    return now.strftime("%a %m%d %I%M%p").upper()

def clean_for_speech(text):
    text = text.replace("#", "")
    emoji_pattern = re.compile(
        r"[\U0001F300-\U0001F5FF"
        r"\U0001F600-\U0001F64F"
        r"\U0001F680-\U0001F6FF"
        r"\U0001F700-\U0001F77F"
        r"\U0001F780-\U0001F7FF"
        r"\U0001F800-\U0001F8FF"
        r"\U0001F900-\U0001F9FF"
        r"\U0001FA00-\U0001FA6F"
        r"\U0001FA70-\U0001FAFF"
        r"\u2600-\u26FF"
        r"\u2700-\u27BF]+", flags=re.UNICODE)
    text = emoji_pattern.sub('', text)
    return text

def trim_emojis_except_numbered(markdown_text):
    emoji_pattern = re.compile(
        r"[\U0001F300-\U0001F5FF"
        r"\U0001F600-\U0001F64F"
        r"\U0001F680-\U0001F6FF"
        r"\U0001F700-\U0001F77F"
        r"\U0001F780-\U0001F7FF"
        r"\U0001F800-\U0001F8FF"
        r"\U0001F900-\U0001F9FF"
        r"\U0001FAD0-\U0001FAD9"
        r"\U0001FA00-\U0001FA6F"
        r"\U0001FA70-\U0001FAFF"
        r"\u2600-\u26FF"
        r"\u2700-\u27BF]+"
    )
    number_pattern = re.compile(r'^\d+\.\s')
    lines = markdown_text.strip().split('\n')
    processed_lines = []
    
    for line in lines:
        if number_pattern.match(line):
            processed_lines.append(line)
        else:
            processed_lines.append(emoji_pattern.sub('', line))
    
    return '\n'.join(processed_lines)

async def generate_audio(text, voice, filename):
    communicate = edge_tts.Communicate(text, voice)
    await communicate.save(filename)
    return filename

def detect_and_convert_links(text):
    # Convert Markdown links [text](url) to HTML <a> tags
    md_link_pattern = re.compile(r'\[(.*?)\]\((https?://[^\s\[\]()<>{}]+)\)')
    text = md_link_pattern.sub(r'<a href="\2" color="blue">\1</a>', text)
    
    # Convert plain URLs to HTML <a> tags, avoiding already tagged links
    url_pattern = re.compile(
        r'(?<!href=")(https?://[^\s<>{}]+)',
        re.IGNORECASE
    )
    text = url_pattern.sub(r'<a href="\1" color="blue">\1</a>', text)
    return text

def apply_emoji_font(text, emoji_font):
    # Protect existing tags
    tag_pattern = re.compile(r'(<[^>]+>)')
    segments = tag_pattern.split(text)
    result = []
    
    # Apply emoji font only to emojis, use DejaVuSans for other text
    emoji_pattern = re.compile(
        r"([\U0001F300-\U0001F5FF"
        r"\U0001F600-\U0001F64F"
        r"\U0001F680-\U0001F6FF"
        r"\U0001F700-\U0001F77F"
        r"\U0001F780-\U0001F7FF"
        r"\U0001F800-\U0001F8FF"
        r"\U0001F900-\U0001F9FF"
        r"\U0001FAD0-\U0001FAD9"
        r"\U0001FA00-\U0001FA6F"
        r"\U0001FA70-\U0001FAFF"
        r"\u2600-\u26FF"
        r"\u2700-\u27BF]+)"
    )
    
    def replace_emoji(match):
        emoji = match.group(1)
        emoji = unicodedata.normalize('NFC', emoji)
        return f'<font face="{emoji_font}">{emoji}</font>'
    
    for segment in segments:
        if tag_pattern.match(segment):
            # Keep tags unchanged
            result.append(segment)
        else:
            # Apply DejaVuSans to non-emoji text, emoji_font to emojis
            parts = []
            last_pos = 0
            for match in emoji_pattern.finditer(segment):
                start, end = match.span()
                if last_pos < start:
                    parts.append(f'<font face="DejaVuSans">{segment[last_pos:start]}</font>')
                parts.append(replace_emoji(match))
                last_pos = end
            if last_pos < len(segment):
                parts.append(f'<font face="DejaVuSans">{segment[last_pos:]}</font>')
            result.append(''.join(parts))
    
    return ''.join(result)

def markdown_to_pdf_content(markdown_text, add_space_before_numbered, headings_to_fonts):
    lines = markdown_text.strip().split('\n')
    pdf_content = []
    number_pattern = re.compile(r'^\d+(\.\d+)*\.\s')
    heading_pattern = re.compile(r'^(#{1,4})\s+(.+)$')
    first_numbered_seen = False
    
    for line in lines:
        line = line.strip()
        if not line:
            continue
            
        if headings_to_fonts and line.startswith('#'):
            heading_match = heading_pattern.match(line)
            if heading_match:
                level = len(heading_match.group(1))
                heading_text = heading_match.group(2).strip()
                formatted_heading = f"<h{level}>{heading_text}</h{level}>"
                pdf_content.append(formatted_heading)
                continue
                
        is_numbered_line = number_pattern.match(line) is not None
        
        if add_space_before_numbered and is_numbered_line:
            if first_numbered_seen and not line.startswith("1."):
                pdf_content.append("")
            if not first_numbered_seen:
                first_numbered_seen = True
        
        line = detect_and_convert_links(line)
        
        # Preserve bold and emphasis formatting
        line = re.sub(r'\*\*(.+?)\*\*', r'<b>\1</b>', line)
        line = re.sub(r'\*([^*]+?)\*', r'<b>\1</b>', line)
                
        pdf_content.append(line)
    total_lines = len(pdf_content)
    return pdf_content, total_lines

def create_pdf(markdown_text, base_font_size, num_columns, add_space_before_numbered, headings_to_fonts, doc_title, longest_line_words, total_lines):
    buffer = io.BytesIO()
    page_width = A4[0] * 2
    page_height = A4[1]
    doc = SimpleDocTemplate(
        buffer,
        pagesize=(page_width, page_height),
        leftMargin=36,
        rightMargin=36,
        topMargin=36,
        bottomMargin=36,
        title=doc_title
    )
    styles = getSampleStyleSheet()
    spacer_height = 10
    pdf_content, total_lines = markdown_to_pdf_content(markdown_text, add_space_before_numbered, headings_to_fonts)
    try:
        available_font_files = glob.glob("*.ttf")
        if not available_font_files:
            st.error("No .ttf font files found.")
            return
        selected_font_path = next((f for f in available_font_files if "NotoEmoji-Bold" in f), None)
        if selected_font_path:
            pdfmetrics.registerFont(TTFont("NotoEmoji-Bold", selected_font_path))
        pdfmetrics.registerFont(TTFont("DejaVuSans", "DejaVuSans.ttf"))
    except Exception as e:
        st.error(f"Font registration error: {e}")
        return
    total_chars = sum(len(line) for line in pdf_content)
    hierarchy_weight = sum(1.5 if line.startswith("<b>") else 1 for line in pdf_content)
    content_density = total_lines * hierarchy_weight + total_chars / 50
    usable_height = page_height - 72 - spacer_height
    usable_width = page_width - 72
    avg_line_chars = total_chars / total_lines if total_lines > 0 else 50
    ideal_lines_per_col = 20
    suggested_columns = max(2, min(4, int(total_lines / ideal_lines_per_col) + 1))
    num_columns = num_columns if num_columns != 0 else suggested_columns
    col_width = usable_width / num_columns
    min_font_size = 6
    max_font_size = 16
    lines_per_col = total_lines / num_columns if num_columns > 0 else total_lines
    target_height_per_line = usable_height / lines_per_col if lines_per_col > 0 else usable_height
    estimated_font_size = int(target_height_per_line / 1.5)
    adjusted_font_size = max(min_font_size, min(max_font_size, estimated_font_size))
    if avg_line_chars > col_width / adjusted_font_size * 10:
        adjusted_font_size = int(col_width / (avg_line_chars / 10))
        adjusted_font_size = max(min_font_size, adjusted_font_size)
    
    # Adjust font size to fit one page based on longest line
    if longest_line_words > 17:
        font_scale = 17 / longest_line_words  # Scale down from reference (17 words)
        adjusted_font_size = max(min_font_size, int(adjusted_font_size * font_scale))
    
    item_style = ParagraphStyle(
        'ItemStyle', parent=styles['Normal'], fontName="DejaVuSans",
        fontSize=adjusted_font_size, leading=adjusted_font_size * 1.15, spaceAfter=1,
        linkUnderline=True
    )
    numbered_bold_style = ParagraphStyle(
        'NumberedBoldStyle', parent=styles['Normal'], fontName="NotoEmoji-Bold",
        fontSize=adjusted_font_size, leading=adjusted_font_size * 1.15, spaceAfter=1,
        linkUnderline=True
    )
    section_style = ParagraphStyle(
        'SectionStyle', parent=styles['Heading2'], fontName="DejaVuSans",
        textColor=colors.darkblue, fontSize=adjusted_font_size * 1.1, leading=adjusted_font_size * 1.32, spaceAfter=2,
        linkUnderline=True
    )
    columns = [[] for _ in range(num_columns)]
    lines_per_column = total_lines / num_columns if num_columns > 0 else total_lines
    current_line_count = 0
    current_column = 0
    number_pattern = re.compile(r'^\d+(\.\d+)*\.\s')
    for item in pdf_content:
        if current_line_count >= lines_per_column and current_column < num_columns - 1:
            current_column += 1
            current_line_count = 0
        columns[current_column].append(item)
        current_line_count += 1
    column_cells = [[] for _ in range(num_columns)]
    for col_idx, column in enumerate(columns):
        for item in column:
            if isinstance(item, str):
                heading_match = re.match(r'<h(\d)>(.*?)</h\1>', item) if headings_to_fonts else None
                if heading_match:
                    level = int(heading_match.group(1))
                    heading_text = heading_match.group(2)
                    heading_style = ParagraphStyle(
                        f'Heading{level}Style', 
                        parent=styles['Heading1'],
                        fontName="DejaVuSans",
                        textColor=colors.darkblue if level == 1 else (colors.black if level > 2 else colors.blue),
                        fontSize=adjusted_font_size * (1.6 - (level-1)*0.15),
                        leading=adjusted_font_size * (1.8 - (level-1)*0.15),
                        spaceAfter=4 - (level-1),
                        spaceBefore=6 - (level-1),
                        linkUnderline=True
                    )
                    column_cells[col_idx].append(Paragraph(apply_emoji_font(heading_text, "NotoEmoji-Bold"), heading_style))
                elif item.startswith("<b>") and item.endswith("</b>"):
                    content = item[3:-4].strip()
                    if number_pattern.match(content):
                        column_cells[col_idx].append(Paragraph(apply_emoji_font(content, "NotoEmoji-Bold"), numbered_bold_style))
                    else:
                        column_cells[col_idx].append(Paragraph(apply_emoji_font(content, "NotoEmoji-Bold"), section_style))
                else:
                    column_cells[col_idx].append(Paragraph(apply_emoji_font(item, "NotoEmoji-Bold"), item_style))
            else:
                column_cells[col_idx].append(Paragraph(apply_emoji_font(str(item), "NotoEmoji-Bold"), item_style))
    max_cells = max(len(cells) for cells in column_cells) if column_cells else 0
    for cells in column_cells:
        cells.extend([Paragraph("", item_style)] * (max_cells - len(cells)))
    table_data = list(zip(*column_cells)) if column_cells else [[]]
    table = Table(table_data, colWidths=[col_width] * num_columns, hAlign='CENTER')
    table.setStyle(TableStyle([
        ('VALIGN', (0, 0), (-1, -1), 'TOP'),
        ('ALIGN', (0, 0), (-1, -1), 'LEFT'),
        ('BACKGROUND', (0, 0), (-1, -1), colors.white),
        ('GRID', (0, 0), (-1, -1), 0, colors.white),
        ('LINEAFTER', (0, 0), (num_columns-1, -1), 0.5, colors.grey),
        ('LEFTPADDING', (0, 0), (-1, -1), 2),
        ('RIGHTPADDING', (0, 0), (-1, -1), 2),
        ('TOPPADDING', (0, 0), (-1, -1), 1),
        ('BOTTOMPADDING', (0, 0), (-1, -1), 1),
    ]))
    story = [Spacer(1, spacer_height), table]
    doc.build(story)
    buffer.seek(0)
    return buffer.getvalue()

def pdf_to_image(pdf_bytes):
    try:
        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
        images = []
        for page in doc:
            pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            images.append(img)
        doc.close()
        return images
    except Exception as e:
        st.error(f"Failed to render PDF preview: {e}")
        return None

# PDF creation and linking functions
WORDS_12 = ["one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten", "eleven", "twelve"]
WORDS_24 = ["one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten",
            "eleven", "twelve", "thirteen", "fourteen", "fifteen", "sixteen", "seventeen", "eighteen", "nineteen", "twenty", "twenty one", "twenty two", "twenty three", "twenty four"]

def create_crossfile_pdfs(source_pdf="TestSource.pdf", target_pdf="TestTarget.pdf"):
    """Create two PDFs with cross-file linking."""
    def create_base_pdf(filename):
        buffer = io.BytesIO()
        c = canvas.Canvas(buffer)
        c.setFont("Helvetica", 12)
        for i, word in enumerate(WORDS_12, 1):
            y = 800 - (i * 20)
            c.drawString(50, y, f"{i}. {word}")
        c.showPage()
        c.save()
        buffer.seek(0)
        with open(filename, "wb") as f:
            f.write(buffer.getvalue())
        buffer.close()

    def add_bookmark_to_seven(pdf_file):
        reader = PdfReader(pdf_file)
        writer = PdfWriter()
        for page in reader.pages:
            writer.add_page(page)
        page = writer.pages[0]
        y_position = 800 - (7 * 20)
        fit = Fit(fit_type="/XYZ", fit_args=[50, y_position, 0])
        writer.add_outline_item("Seven Bookmark", 0, fit=fit)
        with open(pdf_file, "wb") as f:
            writer.write(f)

    def modify_source_pdf(source, target):
        reader = PdfReader(source)
        writer = PdfWriter()
        for page in reader.pages:
            writer.add_page(page)
        buffer = io.BytesIO()
        c = canvas.Canvas(buffer)
        c.setFont("Helvetica", 8)
        seven_y = 800 - (7 * 20)
        c.drawString(90, seven_y - 5, "link")
        c.showPage()
        c.save()
        buffer.seek(0)
        text_pdf = PdfReader(buffer)
        page = writer.pages[0]
        page.merge_page(text_pdf.pages[0])
        link = Link(
            rect=(90, seven_y - 10, 150, seven_y + 10),
            url=f"file://{os.path.abspath(target)}#page=1"
        )
        writer.add_annotation(page_number=0, annotation=link)
        with open(source, "wb") as f:
            writer.write(f)
        buffer.close()

    def add_internal_link(pdf_file):
        reader = PdfReader(pdf_file)
        writer = PdfWriter()
        for page in reader.pages:
            writer.add_page(page)
        one_y = 800 - (1 * 20)
        ten_y = 800 - (10 * 20)
        link = Link(
            rect=(50, one_y - 10, 100, one_y + 10),
            target_page_index=0,
            fit=Fit(fit_type="/XYZ", fit_args=[50, ten_y, 0])
        )
        writer.add_annotation(page_number=0, annotation=link)
        with open(pdf_file, "wb") as f:
            writer.write(f)

    create_base_pdf(source_pdf)
    create_base_pdf(target_pdf)
    add_bookmark_to_seven(target_pdf)
    modify_source_pdf(source, target)
    add_internal_link(source_pdf)
    add_internal_link(target_pdf)
    return source_pdf, target_pdf

def create_selflinking_pdf(pdf_file="SelfLinking.pdf"):
    """Create a PDF with a TOC on page 1 linking to a 1-20 list starting on page 2."""
    buffer = io.BytesIO()
    c = canvas.Canvas(buffer)
    
    # Page 1: Table of Contents
    c.setFont("Helvetica", 14)
    c.drawString(50, 800, "Table of Contents")
    c.setFont("Helvetica", 12)
    toc_y_positions = []
    for i, word in enumerate(WORDS_12, 1):
        y = 760 - (i * 20)
        c.drawString(50, y, f"{word}")
        toc_y_positions.append(y)
    c.showPage()
    
    # Page 2: Numbered list 1-20
    c.setFont("Helvetica", 12)
    list_y_positions = []
    for i, word in enumerate(WORDS_24, 1):
        y = 800 - (i * 20)
        c.drawString(50, y, f"{i}. {word}")
        list_y_positions.append(y)
    c.showPage()
    
    # Save the initial PDF
    c.save()
    buffer.seek(0)
    with open(pdf_file, "wb") as f:
        f.write(buffer.getvalue())
    buffer.close()

    # Add outlines and links
    reader = PdfReader(pdf_file)
    writer = PdfWriter()
    for page in reader.pages:
        writer.add_page(page)

    # Add outline entries
    toc_page = writer.pages[0]
    list_page = writer.pages[1]
    writer.add_outline_item("Table of Contents", 0, fit=Fit(fit_type="/Fit"))
    for i, word in enumerate(WORDS_12, 1):
        y = list_y_positions[i-1]
        writer.add_outline_item(word, 1, fit=Fit(fit_type="/XYZ", fit_args=[50, y, 0]))

    # Add TOC links from page 1 to page 2
    for i, word in enumerate(WORDS_12):
        toc_y = toc_y_positions[i]
        list_y = list_y_positions[i]
        link = Link(
            rect=(50, toc_y - 10, 150, toc_y + 10),
            target_page_index=1,
            fit=Fit(fit_type="/XYZ", fit_args=[50, list_y, 0])
        )
        writer.add_annotation(page_number=0, annotation=link)

    # Save the modified PDF
    with open(pdf_file, "wb") as f:
        writer.write(f)

    return pdf_file

# Streamlit UI
md_files = [f for f in glob.glob("*.md") if os.path.basename(f) != "README.md"]
md_options = [os.path.splitext(os.path.basename(f))[0] for f in md_files]

with st.sidebar:
    st.markdown("### 📄 PDF Options")
    if md_options:
        selected_md = st.selectbox("Select Markdown File", options=md_options, index=0)
        with open(f"{selected_md}.md", "r", encoding="utf-8") as f:
            st.session_state.markdown_content = f.read()
    else:
        st.warning("No markdown file found. Please add one to your folder.")
        selected_md = None
        st.session_state.markdown_content = ""
    
    available_font_files = {os.path.splitext(os.path.basename(f))[0]: f for f in glob.glob("*.ttf")}
    selected_font_name = st.selectbox(
        "Select Emoji Font",
        options=list(available_font_files.keys()),
        index=list(available_font_files.keys()).index("NotoEmoji-Bold") if "NotoEmoji-Bold" in available_font_files else 0
    )
    base_font_size = st.slider("Font Size (points)", min_value=6, max_value=16, value=8, step=1)
    
    add_space_before_numbered = st.checkbox("Add Space Ahead of Numbered Lines", value=False)
    headings_to_fonts = st.checkbox(
        "Headings to Fonts",
        value=False,
        help="Convert Markdown headings (# Heading) to styled fonts"
    )
    auto_columns = st.checkbox("AutoColumns", value=True)
    
    # Calculate document stats
    longest_line_words = 0
    total_lines = 0
    adjusted_font_size_display = base_font_size
    if 'markdown_content' in st.session_state:
        current_markdown = st.session_state.markdown_content
        lines = current_markdown.strip().split('\n')
        total_lines = len([line for line in lines if line.strip()])
        for line in lines:
            if line.strip():
                word_count = len(line.split())
                longest_line_words = max(longest_line_words, word_count)
        if auto_columns:
            if longest_line_words > 38:
                recommended_columns = 2
            elif longest_line_words < 18 and total_lines < 20:
                recommended_columns = 4
            else:
                recommended_columns = 3
        else:
            recommended_columns = 3
        # Adjust font size for one-page fit
        if longest_line_words > 17:
            font_scale = 17 / longest_line_words
            adjusted_font_size_display = max(6, int(base_font_size * font_scale))
        st.markdown("**Document Stats**")
        st.write(f"- Longest Line: {longest_line_words} words")
        st.write(f"- Total Lines: {total_lines}")
        st.write(f"- Recommended Columns: {recommended_columns}")
        st.write(f"- Adjusted Font Size: {adjusted_font_size_display} points")
    
    column_options = [2, 3, 4]
    num_columns = st.selectbox(
        "Number of Columns",
        options=column_options,
        index=column_options.index(recommended_columns)
    )
    st.info("Font size and columns adjust to fit one page.")
    
    st.markdown("### ✍️ Edit Markdown")
    edited_markdown = st.text_area(
        "Input Markdown",
        value=st.session_state.markdown_content,
        height=200,
        key=f"markdown_{selected_md}_{selected_font_name}_{num_columns}"
    )
    
    st.markdown("### 💾 Actions")
    col1, col2 = st.columns(2)
    with col1:
        if st.button("🔄 Update PDF"):
            st.session_state.markdown_content = edited_markdown
            if selected_md:
                with open(f"{selected_md}.md", "w", encoding="utf-8") as f:
                    f.write(edited_markdown)
            st.rerun()
    
    with col2:
        if st.button("✂️ Trim Emojis"):
            trimmed_content = trim_emojis_except_numbered(edited_markdown)
            st.session_state.markdown_content = trimmed_content
            if selected_md:
                with open(f"{selected_md}.md", "w", encoding="utf-8") as f:
                    f.write(trimmed_content)
            st.rerun()
    
    prefix = get_timestamp_prefix()
    st.download_button(
        label="💾 Save Markdown",
        data=st.session_state.markdown_content,
        file_name=f"{prefix} {selected_md}.md" if selected_md else f"{prefix} default.md",
        mime="text/markdown"
    )
    
    st.markdown("### 🔊 Text-to-Speech")
    VOICES = ["en-US-AriaNeural", "en-US-JennyNeural", "en-GB-SoniaNeural", "en-US-GuyNeural", "en-US-AnaNeural"]
    selected_voice = st.selectbox("Select Voice for TTS", options=VOICES, index=0)
    if st.button("Generate Audio"):
        cleaned_text = clean_for_speech(st.session_state.markdown_content)
        audio_filename = f"{prefix} {selected_md} {selected_voice}.mp3" if selected_md else f"{prefix} default {selected_voice}.mp3"
        audio_file = asyncio.run(generate_audio(cleaned_text, selected_voice, audio_filename))
        st.audio(audio_file)
        with open(audio_file, "rb") as f:
            audio_bytes = f.read()
        st.download_button(
            label="💾 Save Audio",
            data=audio_bytes,
            file_name=audio_filename,
            mime="audio/mpeg"
        )
    
    if st.button("📑 Create CrossFile PDFs"):
        with st.spinner("Creating cross-file linked PDFs..."):
            source_pdf, target_pdf = create_crossfile_pdfs()
            st.success(f"Created {source_pdf} and {target_pdf}")
            for pdf_file in [source_pdf, target_pdf]:
                with open(pdf_file, "rb") as f:
                    st.download_button(
                        label=f"💾 Download {pdf_file}",
                        data=f.read(),
                        file_name=pdf_file,
                        mime="application/pdf"
                    )
    
    if st.button("🧪 Create SelfLinking PDF"):
        with st.spinner("Generating self-linking PDF with TOC..."):
            pdf_file = create_selflinking_pdf()
            st.success(f"Generated {pdf_file}")
            with open(pdf_file, "rb") as f:
                pdf_bytes = f.read()
            images = pdf_to_image(pdf_bytes)
            if images:
                st.subheader(f"Preview of {pdf_file}")
                for i, img in enumerate(images):
                    st.image(img, caption=f"{pdf_file} Page {i+1}", use_container_width=True)
            with open(pdf_file, "rb") as f:
                st.download_button(
                    label=f"💾 Download {pdf_file}",
                    data=f.read(),
                    file_name=pdf_file,
                    mime="application/pdf"
                )

with st.spinner("Generating PDF..."):
    pdf_bytes = create_pdf(
        st.session_state.markdown_content,
        base_font_size,
        num_columns,
        add_space_before_numbered,
        headings_to_fonts,
        doc_title=selected_md if selected_md else "Untitled",
        longest_line_words=longest_line_words,
        total_lines=total_lines
    )

with st.container():
    st.markdown("### 📊 PDF Preview")
    pdf_images = pdf_to_image(pdf_bytes)
    if pdf_images:
        for img in pdf_images:
            st.image(img, use_container_width=True)
    else:
        st.info("Download the PDF to view it locally.")

with st.sidebar:
    st.download_button(
        label="💾 Save PDF",
        data=pdf_bytes,
        file_name=f"{prefix} {selected_md}.pdf" if selected_md else f"{prefix} output.pdf",
        mime="application/pdf"
    )