import streamlit as st
from transformers import pipeline
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
import os

# Load Hugging Face Token from Environment (Set this in Hugging Face Spaces Secrets)
HF_TOKEN = os.getenv("HF_TOKEN")

# Select model: Use LLaMA 2 (requires authentication) or Mistral-7B (open-source alternative)
MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"  # Change to "mistralai/Mistral-7B-Instruct-v0.1" if no authentication needed

# Load Model
@st.cache_resource
def load_model():
    try:
        return pipeline("text-generation", model=MODEL_NAME, token=HF_TOKEN)
    except Exception as e:
        st.error(f"❌ Error loading model: {str(e)}")
        return None

generator = load_model()

# Function to Generate Functional Requirement Document
def generate_functional_requirements(topic):
    if generator is None:
        return "Error: Model not loaded properly."

    prompt = f"Generate a comprehensive functional requirement document for {topic} in the banking sector."
    
    # Generate content
    output = generator(prompt, max_length=8000, do_sample=True, temperature=0.8)
    return output[0]['generated_text']

# Function to Save as PDF
def save_to_pdf(content, filename):
    c = canvas.Canvas(filename, pagesize=letter)
    c.setFont("Helvetica", 10)
    
    text = c.beginText(40, 750)
    text.setLeading(14)
    
    for line in content.split("\n"):
        text.textLine(line)
        if text.getY() < 50:  # Start new page when space runs out
            c.drawText(text)
            c.showPage()
            text = c.beginText(40, 750)
            text.setLeading(14)
    
    c.drawText(text)
    c.save()

# Streamlit UI
def main():
    st.title("📄 AI-Powered Functional Requirement Generator for Banking")

    banking_topics = [
        "Core Banking System",
        "Loan Management System",
        "Payment Processing Gateway",
        "Risk and Fraud Detection",
        "Regulatory Compliance Management",
        "Digital Banking APIs",
        "Customer Onboarding & KYC",
        "Treasury Management",
        "Wealth & Portfolio Management"
    ]
    
    topic = st.selectbox("Select a Banking Functional Requirement Topic", banking_topics)
    
    if st.button("Generate Functional Requirement Document"):
        with st.spinner("Generating... This may take a while."):
            content = generate_functional_requirements(topic)
            if "Error" in content:
                st.error(content)
            else:
                filename = "functional_requirement.pdf"
                save_to_pdf(content, filename)
                st.success("✅ Functional Requirement Document Generated!")
                st.download_button(label="📥 Download PDF", data=open(filename, "rb"), file_name=filename, mime="application/pdf")

if __name__ == "__main__":
    main()