import streamlit as st
from docx import Document
import io
from functions import * # get_report_fields, get_report_sections  # Import the necessary functions
import logging

from exa_py import Exa
from groq import Groq
import os

import streamlit as st
from docx import Document
import re
import io
import os
import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.base import MIMEBase
from email import encoders
from email.mime.text import MIMEText
from fpdf import FPDF
from dotenv import load_dotenv
from retrying import retry

# Declare the exa search API
exa = Exa(api_key=os.getenv("EXA_API_KEY"))

# Define your API Model and key
client = Groq(api_key=os.getenv("GROQ_API_KEY"))
utilized_model = "llama3-70b-8192"

# Options for highlights from Exa search
highlights_options = {
    "num_sentences": 7,  # Length of highlights
    "highlights_per_url": 1,  # Get the best highlight for each URL
}

# Initialize session state for tracking progress
if 'session_tracker' not in st.session_state:
    st.session_state.session_tracker = []

# Update session tracker
def update_session_tracker(section_name):
    if section_name not in st.session_state.session_tracker:
        st.session_state.session_tracker.append(section_name)

# Document-related utilities
def sanitize_filename(filename, max_length=50):
    return re.sub(r'[<>:"/\\|?*]', '', filename[:max_length])

def create_document():
    doc = Document()
    doc.add_heading("Generated Report", 0)
    return doc

def add_section_to_doc(doc, section_name, section_content):
    doc.add_heading(section_name, level=1)
    doc.add_paragraph(section_content)
    return doc

def get_docx_bytes(doc):
    doc_io = io.BytesIO()
    doc.save(doc_io)
    doc_io.seek(0)
    return doc_io

# Generic function to dynamically call report generation functions
def call_dynamic_function(function_name, data):
    # Use globals to call the function by name
    function = globals()[function_name]
    return function(data)

# Main info collection and report generation function
def collect_basic_info():
    st.title("Academic Report Generator")

    # Select report type
    report_type = st.selectbox("Select report type", ["PhD Proposal", "Research Paper", "Thesis"])

    # Fetch the dynamic fields and sections from functions.py
    report_fields = get_report_fields(report_type)
    report_sections = get_report_sections(report_type)

    # Create data dictionary by dynamically generating input fields
    data = {field_key: st.text_area(field_label) for field_key, field_label in report_fields}

    # Select section to generate
    section_to_generate = st.selectbox("Choose section to generate", [section[0] for section in report_sections])

    if st.button('Generate Report Section'):
        # Dynamically get the function name and call it
        section_function_name = dict(report_sections)[section_to_generate]
        update_session_tracker(section_to_generate)

        # Generate content and display it
        section_content = call_dynamic_function(section_function_name, data)
        st.subheader(section_to_generate)
        st.write(section_content)

        # Create and download the document
        doc = create_document()
        doc = add_section_to_doc(doc, section_to_generate, section_content)
        doc_bytes = get_docx_bytes(doc)

        sanitized_topic = sanitize_filename(data["research_topic"])
        filename = f"{report_type}_{sanitized_topic}.docx"
        st.download_button(f"Download {section_to_generate} as DOCX", doc_bytes, filename)

# Run the main function
collect_basic_info()