# modules/info_page.py
import streamlit as st
from typing import Dict, List, Optional
from modules.api_utils import fetch_wikipedia_page_details, fetch_wikipedia_images
from io import BytesIO
import base64
import re # Import regex module

def display_page_info(topic: str, num_images: int = 1):
    """Displays detailed information and images for a given topic."""
    st.subheader(f"Information on: {topic}")

    page_details = fetch_wikipedia_page_details(topic)
    
    if page_details:
        st.markdown(f"### {page_details.get('title', topic)}")
        
        with st.expander("Read Full Information", expanded=True):
            extract = page_details.get("extract", "No detailed information found.")
            
            # Attempt to split into paragraphs and identify potential list items
            lines = extract.split('\n')
            formatted_content = []
            for line in lines:
                line = line.strip()
                if not line:
                    continue # Skip empty lines

                # Heuristic for bullet points/list items
                if re.match(r'^[\*\-\d]+\s', line): # Starts with *, -, or digit followed by space
                    formatted_content.append(f"* {line.lstrip('*- ')}") # Format as markdown list
                elif line.endswith(':') or line.endswith('.'): # Likely a full sentence/paragraph
                    formatted_content.append(line)
                else: # Treat as part of previous paragraph or new paragraph
                    if formatted_content and not formatted_content[-1].startswith('*'): # If previous was not a list item
                        formatted_content[-1] += " " + line # Append to last paragraph
                    else:
                        formatted_content.append(line) # New paragraph

            # Join with appropriate spacing for markdown rendering
            final_markdown = ""
            for item in formatted_content:
                if item.startswith('*'):
                    final_markdown += f"{item}\n"
                else:
                    final_markdown += f"{item}\n\n"
            
            st.markdown(final_markdown.strip()) # Render as markdown

        st.markdown("---") # Separator
        
        # Display images
        if num_images > 0:
            st.markdown("#### Images")
            images_found = fetch_wikipedia_images(topic, limit=num_images)
            if images_found:
                # Use st.columns to display images in a grid
                cols = st.columns(min(num_images, len(images_found))) # Adjust columns based on actual images found
                for i, img_url in enumerate(images_found):
                    with cols[i]:
                        try:
                            st.image(img_url, caption=f"Image {i+1}", use_column_width=True)
                        except Exception as e:
                            st.warning(f"Could not load image {i+1}: {e}")
            else:
                st.info("No relevant images found for this topic.")
        
        st.markdown("---") # Separator

        # Download options
        st.markdown("#### Download Information")
        
        text_content = f"Topic: {page_details.get('title', topic)}\n\n{page_details.get('extract', '')}"
        
        # Download as TXT
        st.download_button(
            label="Download as TXT",
            data=text_content,
            file_name=f"{topic}_info.txt",
            mime="text/plain",
            key="download_txt"
        )

        # Download as Markdown (simple version)
        md_content = f"# {page_details.get('title', topic)}\n\n{page_details.get('extract', '')}"
        if images_found:
            for img_url in images_found:
                md_content += f"\n\n![Image]({img_url})"
        
        st.download_button(
            label="Download as Markdown",
            data=md_content,
            file_name=f"{topic}_info.md",
            mime="text/markdown",
            key="download_md"
        )

    else:
        st.error(f"Could not retrieve detailed information for '{topic}'. Please try a different topic.")