import streamlit as st
import base64
import pytesseract
from PIL import Image
from PyPDF2 import PdfReader
from openai import OpenAI
import os

# Retrieve the OpenAI API Key from the secrets
openai_api_key = st.secrets["OPENAI_API_KEY"]
client = OpenAI(api_key=openai_api_key)

# File upload
st.title("Functional Specification Document Processor")
uploaded_file = st.file_uploader("Upload a Functional Spec Document (PDF, TXT, JPG)", type=["pdf", "txt", "jpg", "jpeg"])

# Extract text from different file types
def extract_text(file):
    text = ""
    if file.type == "application/pdf":
        reader = PdfReader(file)
        for page in reader.pages:
            if page and page.extract_text():
                text += page.extract_text()
    elif file.type == "text/plain":
        text = str(file.read(), "utf-8")
    elif file.type in ["image/jpeg", "image/jpg"]:
        image = Image.open(file)
        text = pytesseract.image_to_string(image)
    return text

# Base64 encoding
def encode_base64(text):
    return base64.b64encode(text.encode("utf-8")).decode("utf-8")

# Step 1: Identify epics, features, and user stories in a breadcrumb structure
def identify_structure(processed_text):
    prompt = f"""
    You are a product owner working from a functional specification document. 
    Analyze the document and identify the high-level epics, features, and detailed user stories.
    Structure them in a clear breadcrumb format (Epic > Feature > User Stories). Include titles and short descriptions.
    
    Document:
    {processed_text}
    """
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ],
        max_tokens=2000
    )
    return response.choices[0].message.content.strip()

# Step 2: Generate detailed outcomes for a specific user story
def generate_user_story_details(user_story):
    prompt = f"""
    You are a product owner. Based on the following user story, create a detailed functional breakdown.
    Include:
    - User story description
    - Acceptance criteria
    - Detailed steps and requirements
    
    User Story:
    {user_story}
    """
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ],
        max_tokens=1500
    )
    return response.choices[0].message.content.strip()

if uploaded_file:
    with st.spinner("Processing document..."):
        extracted_text = extract_text(uploaded_file)
        if extracted_text:
            st.subheader("Identified Epics, Features, and User Stories")
            structure = identify_structure(extracted_text)
            st.text_area("Structure", structure, height=300)

            user_story = st.text_input("Copy and paste a user story from the structure above to generate details:")
            if user_story:
                details = generate_user_story_details(user_story)
                st.subheader("Detailed User Story Breakdown")
                st.text_area("Details", details, height=400)

                if st.button("Download as TXT"):
                    st.download_button(label="Download", data=details, file_name="user_story_details.txt", mime="text/plain")
        else:
            st.error("Failed to extract text from the document. Please try again with a clearer document.")