import streamlit as st
import requests
from fpdf import FPDF
import pdfplumber
import os

# Mistral API key (replace with your key or use environment variable)
api_key = os.getenv("MISTRAL_API_KEY", "gz6lDXokxgR6cLY72oomALWcm7vhjRzQ")

# Function to call Mistral API
def call_mistral_api(prompt):
    url = "https://api.mistral.ai/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": "mistral-medium",
        "messages": [
            {"role": "user", "content": prompt}
        ]
    }
    try:
        response = requests.post(url, headers=headers, json=payload)
        response.raise_for_status()  # Raise an error for bad status codes
        return response.json()['choices'][0]['message']['content']
    except requests.exceptions.HTTPError as err:
        return f"HTTP Error: {err}"
    except Exception as err:
        return f"Error: {err}"

# Function to analyze a single requirement
def analyze_requirement(requirement):
    # Detect requirement type
    type_prompt = f"Classify the following requirement as Functional or Non-Functional:\n\n{requirement}\n\nType:"
    req_type = call_mistral_api(type_prompt)

    # Identify stakeholders
    stakeholders_prompt = f"Identify the stakeholders for the following requirement:\n\n{requirement}\n\nStakeholders:"
    stakeholders = call_mistral_api(stakeholders_prompt)

    # Classify domain
    domain_prompt = f"Classify the domain for the following requirement (e.g., Bank, Healthcare, etc.):\n\n{requirement}\n\nDomain:"
    domain = call_mistral_api(domain_prompt)

    # Detect defects
    defects_prompt = f"""Analyze the following requirement and identify ONLY MAJOR defects (e.g., Ambiguity, Incompleteness, etc.). 
    If the requirement is clear and complete, respond with 'No defects.' 
    Requirement: {requirement}
    Defects:"""
    defects = call_mistral_api(defects_prompt)

    # Rewrite requirement
    rewrite_prompt = f"Rewrite the following requirement in a simpler form:\n\n{requirement}\n\nSimplified Requirement:"
    rewritten = call_mistral_api(rewrite_prompt)

    return {
        "Requirement": requirement,
        "Type": req_type,
        "Stakeholders": stakeholders,
        "Domain": domain,
        "Defects": defects,
        "Rewritten": rewritten
    }

# Function to generate a PDF report
def generate_pdf_report(results):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)

    for result in results:
        pdf.cell(200, 10, txt=f"Requirement: {result['Requirement']}", ln=True)
        pdf.cell(200, 10, txt=f"Type: {result['Type']}", ln=True)
        pdf.cell(200, 10, txt=f"Stakeholders: {result['Stakeholders']}", ln=True)
        pdf.cell(200, 10, txt=f"Domain: {result['Domain']}", ln=True)
        pdf.cell(200, 10, txt=f"Defects: {result['Defects']}", ln=True)
        pdf.cell(200, 10, txt=f"Rewritten: {result['Rewritten']}", ln=True)
        pdf.cell(200, 10, txt="-" * 50, ln=True)

    pdf_output = "requirements_report.pdf"
    pdf.output(pdf_output)
    return pdf_output

# Streamlit app
def main():
    st.title("Requirement Analysis Tool")
    st.markdown("**Team Name:** Team Innovators")
    st.markdown("**Model:** Mistral")

    # Input options
    input_option = st.radio("Choose input method:", ("Enter Requirements", "Upload PDF"))

    requirements = []
    if input_option == "Enter Requirements":
        input_text = st.text_area("Enter your requirements (one per line):")
        if input_text:
            requirements = input_text.split("\n")
    else:
        uploaded_file = st.file_uploader("Upload a PDF file:", type="pdf")
        if uploaded_file:
            with pdfplumber.open(uploaded_file) as pdf:
                for page in pdf.pages:
                    requirements.extend(page.extract_text().split("\n"))

    # Analyze requirements
    if st.button("Analyze Requirements"):
        if not requirements:
            st.warning("Please enter or upload requirements.")
        else:
            results = []
            for req in requirements:
                if req.strip():  # Ignore empty lines
                    results.append(analyze_requirement(req.strip()))

            # Display results
            st.subheader("Analysis Results")
            for result in results:
                st.write(f"**Requirement:** {result['Requirement']}")
                st.write(f"**Type:** {result['Type']}")
                st.write(f"**Stakeholders:** {result['Stakeholders']}")
                st.write(f"**Domain:** {result['Domain']}")
                st.write(f"**Defects:** {result['Defects']}")
                st.write(f"**Rewritten:** {result['Rewritten']}")
                st.write("---")

            # Generate and download PDF report
            pdf_report = generate_pdf_report(results)
            with open(pdf_report, "rb") as f:
                st.download_button(
                    label="Download PDF Report",
                    data=f,
                    file_name="requirements_report.pdf",
                    mime="application/pdf"
                )

# Run the app
if __name__ == "__main__":
    main()