import os
import openai
import matplotlib.pyplot as plt
from docx import Document
from docx.shared import Inches, Pt
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
from datetime import datetime
import streamlit as st
from openai import OpenAI

# Initialize OpenAI client
client = OpenAI(api_key=st.secrets.get("OPENAI_API_KEY"))

# Custom CSS for styling
st.markdown("""
<style>
    .stTextInput>div>div>input, .stTextArea>div>div>textarea {
        background-color: #f0f2f6;
        border-radius: 5px;
    }
    .stButton>button {
        background-color: #4CAF50;
        color: white;
        padding: 10px 24px;
        border-radius: 5px;
        border: none;
    }
    .stButton>button:hover {
        background-color: #45a049;
    }
</style>
""", unsafe_allow_html=True)

def init_session_state():
    if 'current_step' not in st.session_state:
        st.session_state.current_step = 1
    if 'user_inputs' not in st.session_state:
        st.session_state.user_inputs = {}

def generate_text(prompt, role="user"):
    try:
        response = client.chat.completions.create(
            model="gpt-4-1106-preview",
            messages=[
                {"role": "system", "content": "You are an expert banking analyst."},
                {"role": role, "content": prompt}
            ]
        )
        return response.choices[0].message.content
    except Exception as e:
        st.error(f"Error generating text: {e}")
        return "(*Content generation failed. Please fill manually.*)"

def create_fsd_document(inputs):
    doc = Document()
    
    # Set default font
    style = doc.styles['Normal']
    font = style.font
    font.name = 'Calibri'
    font.size = Pt(11)

    # Title Page
    doc.add_heading('Functional Specification Document (FSD)', 0)
    doc.add_paragraph("Project Name: " + inputs.get('project_name', 'Banking System'))
    doc.add_paragraph(f"Version: {inputs.get('version', '1.0')}")
    doc.add_paragraph(f"Prepared By: {inputs.get('author', 'Your Name')}")
    doc.add_paragraph(f"Date: {datetime.now().strftime('%d/%m/%Y')}")
    doc.add_page_break()

    # Table of Contents
    doc.add_heading("Table of Contents", 1)
    toc = [
        "1. Introduction",
        "2. System Overview",
        "3. Functional Requirements",
        "4. Non-Functional Requirements",
        "5. Security & Compliance",
        "6. System Architecture",
        "7. Process Flows",
        "8. Data Requirements",
        "9. User Interface",
        "10. Reporting",
        "11. Error Handling",
        "12. Assumptions",
        "13. Dependencies",
        "14. Glossary"
    ]
    for item in toc:
        doc.add_paragraph(item, style='ListBullet')
    doc.add_page_break()

    # Detailed Sections
    sections = {
        "1. Introduction": f"Write a comprehensive introduction for {inputs['project_name']} covering: {inputs['brief_description']}",
        "2. System Overview": f"Create detailed system overview for {inputs['project_name']} including: {inputs['system_components']}",
        "3. Functional Requirements": f"Expand these functional requirements into detailed specifications: {inputs['functional_reqs']}",
        "4. Non-Functional Requirements": f"Create non-functional requirements for: {inputs['non_functional_reqs']}",
        "5. Security & Compliance": f"Detail security measures and compliance requirements for: {inputs['security_needs']}",
    }

    for header, prompt in sections.items():
        doc.add_heading(header, level=1)
        content = generate_text(prompt)
        for paragraph in content.split('\n'):
            if paragraph.strip():
                p = doc.add_paragraph(paragraph)
                p.paragraph_format.space_after = Pt(6)

        # Add detailed subsections
        if header == "3. Functional Requirements":
            doc.add_heading("3.1 Detailed Use Cases", level=2)
            use_cases = generate_text(f"Create detailed use cases for: {inputs['functional_reqs']}")
            doc.add_paragraph(use_cases)

        doc.add_page_break()

    # Add diagrams
    doc.add_heading("6. System Architecture", 1)
    fig = create_architecture_diagram()
    diagram_path = "system_architecture.png"
    fig.savefig(diagram_path)
    doc.add_picture(diagram_path, width=Inches(6))
    
    return doc

def create_architecture_diagram():
    fig, ax = plt.subplots(figsize=(10, 6))
    components = [
        ("User Interface", 0.5, 0.9),
        ("Application Layer", 0.3, 0.7),
        ("Business Logic", 0.7, 0.7),
        ("Database", 0.5, 0.5),
        ("External APIs", 0.5, 0.3)
    ]
    
    for text, x, y in components:
        ax.text(x, y, text, 
                ha='center', va='center',
                bbox=dict(facecolor='lightblue', edgecolor='black', boxstyle='round,pad=0.5'),
                fontsize=10)
    
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.axis('off')
    return fig

def main():
    init_session_state()
    st.title("FSD Document Generator")
    st.subheader("Create Custom Functional Specification Documents")

    if st.session_state.current_step == 1:
        with st.form("basic_info"):
            st.header("Step 1: Basic Information")
            st.session_state.user_inputs['project_name'] = st.text_input("Project Name")
            st.session_state.user_inputs['brief_description'] = st.text_area("Brief Description", height=150)
            if st.form_submit_button("Next →"):
                st.session_state.current_step = 2
                st.rerun()

    elif st.session_state.current_step == 2:
        with st.form("detailed_info"):
            st.header("Step 2: Detailed Requirements")
            st.session_state.user_inputs['functional_reqs'] = st.text_area(
                "Functional Requirements (comma-separated)",
                height=100,
                value="User authentication, Transaction monitoring, Report generation"
            )
            st.session_state.user_inputs['non_functional_reqs'] = st.text_input(
                "Non-Functional Requirements",
                value="High availability, GDPR compliance"
            )
            st.session_state.user_inputs['security_needs'] = st.text_input(
                "Security Requirements",
                value="Data encryption, Role-based access control"
            )
            if st.form_submit_button("Generate FSD"):
                st.session_state.current_step = 3
                st.rerun()

    elif st.session_state.current_step == 3:
        st.header("Step 3: Generated Document")
        with st.spinner("Generating 10+ page FSD document..."):
            doc = create_fsd_document(st.session_state.user_inputs)
            output_folder = "generated_docs"
            os.makedirs(output_folder, exist_ok=True)
            doc_path = os.path.join(output_folder, "FSD_Document.docx")
            doc.save(doc_path)
            
            with open(doc_path, "rb") as file:
                st.download_button(
                    label="Download FSD Document",
                    data=file,
                    file_name="FSD_Document.docx",
                    mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document"
                )
            
            if st.button("Start New Document"):
                st.session_state.current_step = 1
                st.session_state.user_inputs = {}
                st.rerun()

if __name__ == "__main__":
    main()