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):
    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}

    # Add an option to generate all sections
    section_to_generate = st.selectbox("Choose section to generate", [section[0] for section in report_sections])

    if section_to_generate == "All Sections":
        if st.button('Generate All Sections'):
            # Generate all sections content at once
            all_sections_content = generate_all_sections(data)

            # Create a new document for all sections combined
            combined_doc = create_document()

            for section_name, section_content in all_sections_content: #.items():
               st.subheader(section_name)
               st.write(section_content)  # Display the generated content

                # Add each section's content to the combined document
               combined_doc = add_section_to_doc(combined_doc, section_name, section_content)

            # Get bytes for the combined document and create a download button
            combined_doc_bytes = get_docx_bytes(combined_doc)
            sanitized_topic = sanitize_filename(data["research_topic"])
            filename = f"{report_type}_{sanitized_topic}_all_sections.docx"
            st.download_button(f"Download All Sections as DOCX", combined_doc_bytes, filename)

    if section_to_generate == "All Sections_0":
        for section_name, generate_prompt_func in sections_to_process:
            prompt = generate_prompt_func(data)
            section_content = call_llm(prompt)
            st.subheader(section_name)
            st.write(section_content)
            
            # Update document and create download link
            doc = add_section_to_doc(doc, section_name, section_content)
            doc_bytes = get_docx_bytes(doc)

            st.download_button(
                label=f"Download {section_name} as DOCX",
                data=doc_bytes,
                file_name=f"{section_name.replace(' ', '_').lower()}.docx",
                mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document"
            )

            # Save section content in a file
            #filename = f"{company_name}_{section_name}.docx"
            # Sanitize section name and limit file name length
            sanitized_section_name = sanitize_filename(section_name, max_length=50)
            
            # Generate final file name and ensure it is not too long
            filename = f"{sanitized_company_name}_{sanitized_section_name}.docx"
            
            with open(filename, 'wb') as f:
                f.write(doc_bytes.getbuffer())

            # Send email with section content and attachment
            #send_email_with_attachment(email, f"{section_name} of Your Business Proposal", "Please find attached the section of your business proposal.", filename, section_content)

    
    else:
        if st.button('Generate Report Section'):
            # Dynamically get the function name and call it for individual sections
            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)  # Display the generated content

            # Create and download the document for this specific section only
            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()