import os
import streamlit as st
from dotenv import load_dotenv
from langchain_core.messages import AIMessage, HumanMessage
from langchain_community.llms import HuggingFaceEndpoint
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

load_dotenv()

api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")

# Define the repository ID and task
repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
task = "text-generation"

st.set_page_config(page_title="CoDev.AI",page_icon= "🖥️")
st.title("CoDev.AI 🤖")

template = """
You are a software developer expert and software architect chatbot your name is CoDev.AI designed to help users to generate code, command, refactor code, debug issue, fix the issue, generate the commments for given code snippet, check the designs, programmeing paradigm, architect releated information. Here are some scenarios you should be able to handle:

1. Generate code: Assist users with generating the code for the given requirement, generate code by default in python language or the user asked coding language

2. Generate commands: Help users to the linux shell commands.

3. Generate comments: Generate the comments for the code snippet, where all the arguments and return are well explained.

4. Explain error: Help the users to understand the error stack and debug the issue and find root cause.

5. Coding Tips: Offer the user with some coding tips which helps code readbility and following the SOLID principles.

6. Code Optimization: Help the users to optimize the code to improve the performance, improve the space and time complexity

7. Code Explanation: If the code is given, identify which language and give explanation for each line for the purposes

8. Code Formatter: Format the code, identify the coding language and always format and return response prettier

9. SQL Optimisation: Generate query for the given requirement in required database.

Always give explanation for each line or steps, but user asked for command only, then generate command only.

Always give multiple solution and suggest the best from that

Please ensure responses are informative, accurate, and tailored to the user's queries and preferences. Use natural language to engage users and provide a seamless experience throughout their coding and learning.

Chat history:
{chat_history}

User question:
{user_question}
"""

prompt = ChatPromptTemplate.from_template(template)

def get_response(user_query, chat_history):
    # Initialize the Hugging Face Endpoint
    llm = HuggingFaceEndpoint(
        huggingfacehub_api_token=api_token,
        repo_id=repo_id,
        task=task
    )

    chain = prompt | llm | StrOutputParser()

    response = chain.invoke({
        "chat_history": chat_history,
        "user_question": user_query,
    })

    return response

# Initialize session state
if "chat_history" not in st.session_state:
    st.session_state.chat_history = [
        AIMessage(content="Hello, I am CoDev.AI How can I help you?"),
    ]

# Display chat history
for message in st.session_state.chat_history:
    if isinstance(message, AIMessage):
        with st.chat_message("AI"):
            st.write(message.content)
    elif isinstance(message, HumanMessage):
        with st.chat_message("Human"):
            st.write(message.content)

# User input
user_query = st.chat_input("Type your message here...")
if user_query is not None and user_query != "":
    st.session_state.chat_history.append(HumanMessage(content=user_query))

    with st.chat_message("Human"):
        st.markdown(user_query)

    response = get_response(user_query, st.session_state.chat_history)

    # Remove any unwanted prefixes from the response
    response = response.replace("AI response:", "").replace("chat response:", "").replace("bot response:", "").strip()

    with st.chat_message("AI"):
        st.write(response)

    st.session_state.chat_history.append(AIMessage(content=response))