import os
import csv
import streamlit as st
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_huggingface import HuggingFaceEndpoint
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain, RetrievalQA
from huggingface_hub import login

# Login to Hugging Face
login(token=st.secrets["HF_TOKEN"])

# Load FAISS index and ensure it only happens once
if 'db' not in st.session_state:
    st.session_state.db = FAISS.load_local(
        "faiss_index", 
        HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L12-v2'),
        allow_dangerous_deserialization=True
    )

# Use session state for retriever
retriever = st.session_state.db.as_retriever(
    search_type="mmr",
    search_kwargs={'k': 1}
)

# Define prompt template
prompt_template = """
### [INST]
Instruction: You are a Q&A assistant. Your goal is to answer questions as accurately as possible based on the instructions and context provided without using prior knowledge. You answer in FRENCH
Analyse carefully the context and provide a direct answer based on the context. If the user said Bonjour or Hello your only answer will be Hi! comment puis-je vous aider?
Answer in french only

{context}
Vous devez répondre aux questions en français.
### QUESTION:
{question}
[/INST]
Answer in french only
Vous devez répondre aux questions en français.
"""

repo_id = "mistralai/Mistral-7B-Instruct-v0.3"

# Load the model only once
if 'mistral_llm' not in st.session_state:
    st.session_state.mistral_llm = HuggingFaceEndpoint(
        repo_id=repo_id, 
        max_length=2048, 
        temperature=0.05, 
        huggingfacehub_api_token=st.secrets["HF_TOKEN"]
    )

# Create prompt and LLM chain
prompt = PromptTemplate(
    input_variables=["question"],
    template=prompt_template,
)
llm_chain = LLMChain(llm=st.session_state.mistral_llm, prompt=prompt)

# Create QA chain
qa = RetrievalQA.from_chain_type(
    llm=st.session_state.mistral_llm,
    chain_type="stuff",
    retriever=retriever,
    chain_type_kwargs={"prompt": prompt},
)

import psycopg2
import streamlit as st
from datetime import datetime
import os

# PostgreSQL connection setup using secrets from Hugging Face Spaces
def create_connection():
    conn = psycopg2.connect(
        host=os.getenv("DB_HOST"),
        database=os.getenv("DB_NAME"),
        user=os.getenv("DB_USER"),
        password=os.getenv("DB_PASSWORD"),
        port=os.getenv("DB_PORT")
    )
    return conn

def create_table(conn):
    with conn.cursor() as cur:
        cur.execute('''
        CREATE TABLE IF NOT EXISTS feedback (
            id SERIAL PRIMARY KEY,
            timestamp TIMESTAMP NOT NULL,
            rating INTEGER NOT NULL,
            comment TEXT NOT NULL
        );
        ''')
        conn.commit()

# Streamlit interface with improved aesthetics
st.set_page_config(page_title="Alter-IA Chat", page_icon="🤖")

# Define function to handle user input and display chatbot response
def chatbot_response(user_input):
    response = qa.run(user_input)
    return response

# Create columns for logos
col1, col2, col3 = st.columns([2, 3, 2])

with col1:
    st.image("Design 3_22.png", width=150, use_column_width=True)

with col3:
    st.image("Altereo logo 2023 original - eau et territoires durables.png", width=150, use_column_width=True)

# CSS for styling
st.markdown("""
    <style>
    .centered-text {
        text-align: center;
    }
    .centered-orange-text {
        text-align: center;
        color: darkorange;
    }
    .star-rating {
        display: flex;
        flex-direction: row-reverse;
        justify-content: center;
        cursor: pointer;
    }
    .star-rating input[type="radio"] {
        display: none;
    }
    .star-rating label {
        font-size: 2em;
        color: #ddd;
        padding: 0 5px;
        transition: color 0.3s;
    }
    .star-rating input[type="radio"]:checked ~ label {
        color: gold;
    }
    .star-rating input[type="radio"]:hover ~ label {
        color: gold;
    }
    </style>
    """, unsafe_allow_html=True)

st.markdown('<h3 class="centered-text">🤖 AlteriaChat 🤖 </h3>', unsafe_allow_html=True)
st.markdown('<p class="centered-orange-text">"Votre Réponse à Chaque Défi Méthodologique "</p>', unsafe_allow_html=True)

# Input and button for user interaction
user_input = st.text_input("You:", "")
submit_button = st.button("Ask 📨")

if submit_button:
    if user_input.strip() != "":
        bot_response = chatbot_response(user_input)
        st.markdown("### Bot:")
        st.text_area("", value=bot_response, height=300)

        # Add rating and comment section
        st.markdown("---")
        st.markdown("#### Rate the Response:")
        
        # Custom star rating HTML
        rating_html = """
        <div class="star-rating">
            <input type="radio" id="5-stars" name="rating" value="5"><label for="5-stars">★</label>
            <input type="radio" id="4-stars" name="rating" value="4"><label for="4-stars">★</label>
            <input type="radio" id="3-stars" name="rating" value="3" checked><label for="3-stars">★</label>
            <input type="radio" id="2-stars" name="rating" value="2"><label for="2-stars">★</label>
            <input type="radio" id="1-star" name="rating" value="1"><label for="1-star">★</label>
        </div>
        """
        st.markdown(rating_html, unsafe_allow_html=True)
        
        # Get the selected rating via JavaScript
        rating = st.text_input("Selected Rating:", value="3", key="rating_input", label_visibility="hidden")

        comment = st.text_area("Your Comment:")

        # Submit feedback
        feedback_button = st.button("Submit Feedback")

        if feedback_button:
            if comment.strip() == "":
                st.warning("⚠ Please provide a comment.")
            else:
                st.success("Thank you for your feedback!")
                
                # Store feedback in PostgreSQL
                conn = create_connection()
                with conn.cursor() as cur:
                    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    cur.execute('INSERT INTO feedback (timestamp, rating, comment) VALUES (%s, %s, %s)', 
                                (timestamp, int(rating), comment))
                    conn.commit()
                conn.close()

    else:
        st.warning("⚠ Please enter a message.")

# Motivational quote at the bottom
st.markdown("---")
st.markdown("La collaboration est la clé du succès. Chaque question trouve sa réponse, chaque défi devient une opportunité.")