import streamlit as st
import requests
from fpdf import FPDF
import os
import time
from datetime import datetime
import groq

# API keys (replace with your keys or use environment variables)
mistral_api_key = os.getenv("MISTRAL_API_KEY", "your_mistral_api_key")
groq_api_key = os.getenv("GROQ_API_KEY", "your_groq_api_key")

# Initialize Groq client
groq_client = groq.Client(api_key=groq_api_key)

# Function to call Mistral API
def call_mistral_api(prompt):
    url = "https://api.mistral.ai/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {mistral_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()
        return response.json()['choices'][0]['message']['content']
    except requests.exceptions.HTTPError as err:
        if response.status_code == 429:
            st.warning("Rate limit exceeded. Please wait a few seconds and try again.")
            time.sleep(5)
            return call_mistral_api(prompt)
        return f"HTTP Error: {err}"
    except Exception as err:
        return f"Error: {err}"

# Function to call Groq API
def call_groq_api(prompt):
    try:
        response = groq_client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[{"role": "user", "content": prompt}]
        )
        return response.choices[0].message.content
    except Exception as err:
        st.error(f"Error: {err}")
        return f"Error: {err}"

# Streamlit UI
def main():
    st.title("AI Powered Requirement Analysis")
    st.markdown("**Team Name:** Sadia, Areeba, Rabbia, Tesmia")
    st.markdown("**Models:** Mistral (Classification & Domain) + Groq (Defects & Rewriting)")
    
    input_text = st.text_area("Enter requirements (one per line or separated by periods):")
    requirements = [req.strip() for req in input_text.replace("\n", ".").split(".") if req.strip()]
    
    if st.button("Analyze Requirements"):
        if not requirements:
            st.warning("Please enter requirements.")
        else:
            results = []
            for req in requirements:
                if req.strip():
                    results.append(analyze_requirement(req.strip()))
            
            st.subheader("Analysis Results")
            for i, result in enumerate(results, start=1):
                st.write(f"### Requirement R{i}: {result['Requirement']}")
                st.write(f"**Type:** {result['Type']}")
                st.write(f"**Domain:** {result['Domain']}")
                st.write(f"**Defects:** {result['Defects']}")
                st.write(f"**Rewritten:** {result['Rewritten']}")
                st.write("---")

if __name__ == "__main__":
    main()