import streamlit as st
import json
import os
from transformers import pipeline
from langchain.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings

# ====== LOGIN + SIGNUP SYSTEM START ======
USER_FILE = "users.json"

# Load users from file
def load_users():
    if os.path.exists(USER_FILE):
        with open(USER_FILE, "r") as f:
            return json.load(f)
    return {}

# Save user to file
def save_users(users):
    with open(USER_FILE, "w") as f:
        json.dump(users, f)

users = load_users()

# Manage session
if "logged_in" not in st.session_state:
    st.session_state.logged_in = False
if "current_user" not in st.session_state:
    st.session_state.current_user = None

def login_form():
    st.subheader("🔐 Login")
    username = st.text_input("Username")
    password = st.text_input("Password", type="password")
    if st.button("Login"):
        if username in users and users[username] == password:
            st.session_state.logged_in = True
            st.session_state.current_user = username
            st.success("✅ Login successful!")
            st.rerun()
        else:
            st.error("❌ Invalid credentials")

def signup_form():
    st.subheader("📝 Sign Up")
    new_user = st.text_input("New Username")
    new_pass = st.text_input("New Password", type="password")
    if st.button("Create Account"):
        if new_user in users:
            st.warning("⚠️ User already exists")
        else:
            users[new_user] = new_pass
            save_users(users)
            st.success("✅ Account created! Please login.")
            st.rerun()

# Show login/signup
if not st.session_state.logged_in:
    st.title("Welcome to the GRC Chatbot Portal")
    option = st.radio("Select an option", ["Login", "Sign Up"])
    if option == "Login":
        login_form()
    else:
        signup_form()
    st.stop()

# ====== LOGIN SYSTEM END ======

# ====== Chatbot Interface ======
st.set_page_config(page_title="GRC Chatbot", page_icon="🛡️", layout="centered")
st.markdown("""
    <style>
    .main {
        background-color: #0F1117;
        color: white;
        font-family: 'Segoe UI', sans-serif;
    }
    .block-container {
        padding-top: 2rem;
        padding-bottom: 2rem;
    }
    .stTextInput>div>div>input {
        background-color: #1c1e26;
        color: white;
    }
    .stTextInput label, .stTextArea label {
        color: #ffffff;
    }
    .chat-box {
        background-color: #1c1e26;
        padding: 15px;
        margin: 10px 0;
        border-radius: 8px;
        border: 1px solid #303030;
    }
    </style>
""", unsafe_allow_html=True)

if "history" not in st.session_state:
    st.session_state.history = []

# Add logout button
if st.button("🚪 Logout"):
    st.session_state.logged_in = False
    st.session_state.current_user = None
    st.rerun()

# Load summarizer
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

# Load documents
docs = []
for file in ["cyber_docs/nist.txt", "cyber_docs/iso27001.txt", "cyber_docs/gdpr.txt"]:
    loader = TextLoader(file)
    docs.extend(loader.load())

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
split_docs = text_splitter.split_documents(docs)

embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectorstore = FAISS.from_documents(split_docs, embeddings)

def simple_llm(query):
    matched_docs = vectorstore.similarity_search(query, k=5)
    keyword = ""
    if "gdpr" in query.lower():
        keyword = "gdpr"
    elif "nist" in query.lower():
        keyword = "nist"
    elif "iso" in query.lower() or "27001" in query:
        keyword = "iso"

    filtered_docs = [doc for doc in matched_docs if keyword in doc.metadata['source'].lower()]
    if not filtered_docs:
        filtered_docs = matched_docs

    combined_text = " ".join([doc.page_content for doc in filtered_docs])
    summary = summarizer(combined_text, max_length=200, min_length=50, do_sample=False)
    return summary[0]['summary_text']

# UI
st.title("🛡️ AI Chatbot for GRC")
uploaded_file = st.file_uploader("📁 Upload a new GRC .txt file", type="txt")
if uploaded_file is not None:
    os.makedirs("cyber_docs", exist_ok=True)
    file_path = os.path.join("cyber_docs", uploaded_file.name)
    with open(file_path, "wb") as f:
        f.write(uploaded_file.getbuffer())
    st.success(f"✅ {uploaded_file.name} uploaded successfully!")

    loader = TextLoader(file_path)
    new_docs = loader.load()
    docs.extend(new_docs)
    split_docs = text_splitter.split_documents(docs)
    vectorstore = FAISS.from_documents(split_docs, embeddings)

query = st.text_input("Ask your GRC question:")

if st.button("🧹 Clear Chat"):
    st.session_state.history = []
    st.rerun()

st.markdown("""
**💡 Example Queries:**  
- What are the functions of NIST?  
- Explain GDPR principles.  
- What is ISO 27001 risk assessment?
""")
if query:
    try:
        result = simple_llm(query)
        st.session_state.history.append((query, result))
        st.subheader("🔍 Answer:")
        st.write(result)

        if st.session_state.history:
            st.markdown("---")
            st.subheader("💬 Chat History")
            for q, a in reversed(st.session_state.history):
                st.markdown(f"""
<div class='chat-box'>
    <p><b>🧑 You:</b> {q}</p>
    <p><b>🤖 Bot:</b> {a}</p>
</div>
""", unsafe_allow_html=True)

    except Exception as e:
        st.error(f"An error occurred: {e}")
else:
    st.info("Please enter a query related to GRC (NIST, ISO 27001, GDPR).")