import os
from langchain_huggingface import HuggingFaceEndpoint
from transformers import pipeline
import streamlit as st
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_community.document_loaders import YoutubeLoader, WebBaseLoader

model_id="unsloth/Llama-3.2-1B-Instruct"

model = pipeline(
    "text-generation",
    model=model_id,
    torch_dtype="auto",
    token=os.getenv("HF_TOKEN")
)

url = "https://www.youtube.com/watch?v=QsYGlZkevEg"
loader = YoutubeLoader.from_youtube_url(url, add_video_info=False)

data = loader.load()

def generate(prompt: str) -> str:
    ## 1. Transform `input` into a desired format (e.g. it can be simply a string, or a list of dictionaries)
    ## The format of the input depends on the model you are using. You should check the model's documentation.
    input = prompt

    response = model(input, max_new_tokens=512)

    # 2. Make sure to return just the content of the AI response (most of the time the model returns a dictionary with additional information)
    return response[0]['generated_text']

generate("Hello World!")

BASE_PROMPT = """
You are provided with a context and a user query. Your task is to answer the user query based on the provided context.

## Context
{data}

## Query
{query}

## Answer
"""


# model_id="mistralai/Mistral-7B-Instruct-v0.3"
# def get_llm_hf_inference(model_id=model_id, max_new_tokens=128, temperature=0.1, top_k=50, top_p=0.9):
#     """
#     Returns a language model for HuggingFace inference.

#     Parameters:
#     - model_id (str): The ID of the HuggingFace model repository.
#     - max_new_tokens (int): The maximum number of new tokens to generate.
#     - temperature (float): The temperature for sampling from the model.
#     - top_k (int): The number of highest probability tokens to consider.
#     - top_p (float): The cumulative probability threshold for token selection.

#     Returns:
#     - llm (HuggingFaceEndpoint): The language model for HuggingFace inference.
#     """
#     llm = HuggingFaceEndpoint(
#         repo_id=model_id,
#         max_new_tokens=max_new_tokens,
#         temperature=temperature,
#         top_k=top_k,
#         top_p=top_p,
#         token=os.getenv("HF_TOKEN")
#     )
#     return llm

# # Configure the Streamlit app
# st.set_page_config(page_title="HuggingFace ChatBot", page_icon="🤗")
# st.title("Personal HuggingFace ChatBot")
# st.markdown(f"*This is a simple chatbot that uses the HuggingFace transformers library to generate responses to your text input. It uses the {model_id}.*")

# # Initialize session state for avatars
# if "avatars" not in st.session_state:
#     st.session_state.avatars = {'user': None, 'assistant': None}

# # Initialize session state for user text input
# if 'user_text' not in st.session_state:
#     st.session_state.user_text = None

# # Initialize session state for model parameters
# if "max_response_length" not in st.session_state:
#     st.session_state.max_response_length = 256

# if "top_k" not in st.session_state:
#     st.session_state.top_k = 1.0

# if "top_p" not in st.session_state:
#     st.session_state.top_p = 0.95

# if "temperature" not in st.session_state:
#     st.session_state.temperature = 1.0

# if "system_message" not in st.session_state:
#     st.session_state.system_message = "friendly AI conversing with a human user"

# if "starter_message" not in st.session_state:
#     st.session_state.starter_message = "Hello, there! How can I help you today?"
    
    
# # Sidebar for settings
# with st.sidebar:
#     st.header("System Settings")

#     # AI Settings
#     st.session_state.system_message = st.text_area(
#         "System Message", value="You are a friendly AI conversing with a human user."
#     )
#     st.session_state.starter_message = st.text_area(
#         'First AI Message', value="Hello, there! How can I help you today?"
#     )

#     # Model Settings
#     st.session_state.max_response_length = st.number_input(
#         "Max Response Length", value=128
#     )


#     selected_option = st.selectbox("Choose parametrs", ("Pricise", "Balanced", "Creative", "Custom"))

#     if selected_option == "Precise":
#         st.session_state.temperature = 0.1
#         st.session_state.top_k = 10.0
#         st.session_state.top_p = 0.8
#     elif selected_option == "Balanced":
#         st.session_state.temperature = 1.0
#         st.session_state.top_k = 100.0
#         st.session_state.top_p = 0.9
#     elif selected_option == "Creative":
#         st.session_state.temperature = 10.0
#         st.session_state.top_k = 1500.0
#         st.session_state.top_p = 0.99
#     elif selected_option == "Custom":
#         st.session_state.temperature = st.number_input(
#     "Temperature", value=st.session_state.temperature, min_value=0.0
#     )
#         st.session_state.top_k = st.number_input(
#         "Top sK", value=st.session_state.top_k, min_value=1.0, step=1.0
#     )
#         st.session_state.top_p = st.number_input(
#         "Top sP", value=st.session_state.top_p, min_value=0.0, max_value=1.0
#     )
    
#     # Avatar Selection
#     st.markdown("*Select Avatars:*")
#     col1, col2 = st.columns(2)
#     with col1:
#         st.session_state.avatars['assistant'] = st.selectbox(
#             "AI Avatar", options=["🤗", "💬", "🤖"], index=0
#         )
#     with col2:
#         st.session_state.avatars['user'] = st.selectbox(
#             "User Avatar", options=["👤", "👱‍♂️", "👨🏾", "👩", "👧🏾"], index=0
#         )
#     # Reset Chat History
#     reset_history = st.button("Reset Chat History")
    
# # Initialize or reset chat history
# if "chat_history" not in st.session_state or reset_history:
#     st.session_state.chat_history = [{"role": "assistant", "content": st.session_state.starter_message}]

# def get_response(system_message, chat_history, user_text, 
#                  eos_token_id=['User'], max_new_tokens=256, get_llm_hf_kws={}):
#     """
#     Generates a response from the chatbot model.

#     Args:
#         system_message (str): The system message for the conversation.
#         chat_history (list): The list of previous chat messages.
#         user_text (str): The user's input text.
#         eos_token_id (list, optional): The list of end-of-sentence token IDs.
#         max_new_tokens (int, optional): The maximum number of new tokens to generate.
#         get_llm_hf_kws (dict, optional): Additional keyword arguments for the get_llm_hf function.

#     Returns:
#         tuple: A tuple containing the generated response and the updated chat history.
#     """
#     # Set up the model
#     hf = get_llm_hf_inference(
#         model_id=model_id,
#         max_new_tokens=max_new_tokens,
#         temperature=st.session_state.temperature,
#         top_k=st.session_state.top_k,
#         top_p=st.session_state.top_p
#     )

#     # Create the prompt template
#     prompt = PromptTemplate.from_template(
#         (
#             "[INST] {system_message}"
#             "\nCurrent Conversation:\n{chat_history}\n\n"
#             "\nUser: {user_text}.\n [/INST]"
#             "\nAI:"
#         )
#     )
#     # Make the chain and bind the prompt
#     chat = prompt | hf.bind(skip_prompt=True) | StrOutputParser(output_key='content')

#     # Generate the response
#     response = chat.invoke(input=dict(system_message=system_message, user_text=user_text, chat_history=chat_history))
#     response = response.split("AI:")[-1]

#     # Update the chat history
#     chat_history.append({'role': 'user', 'content': user_text})
#     chat_history.append({'role': 'assistant', 'content': response})
#     return response, chat_history


# # Chat interface
# chat_interface = st.container(border=True)
# with chat_interface:
#     output_container = st.container()
#     st.session_state.user_text = st.chat_input(placeholder="Enter your text here.")
    
# # Display chat messages
# with output_container:
#     # For every message in the history
#     for message in st.session_state.chat_history:
#         # Skip the system message
#         if message['role'] == 'system':
#             continue
            
#         # Display the chat message using the correct avatar
#         with st.chat_message(message['role'], 
#                              avatar=st.session_state['avatars'][message['role']]):
#             st.markdown(message['content'])
            
#  # When the user enter new text:
#     if st.session_state.user_text:
        
#         # Display the user's new message immediately
#         with st.chat_message("user", 
#                              avatar=st.session_state.avatars['user']):
#             st.markdown(st.session_state.user_text)
            
#         # Display a spinner status bar while waiting for the response
#         with st.chat_message("assistant", 
#                              avatar=st.session_state.avatars['assistant']):

#             with st.spinner("Thinking..."):
#                 # Call the Inference API with the system_prompt, user text, and history
#                 response, st.session_state.chat_history = get_response(
#                     system_message=st.session_state.system_message, 
#                     user_text=st.session_state.user_text,
#                     chat_history=st.session_state.chat_history,
#                     max_new_tokens=st.session_state.max_response_length,
#                 )
#                 st.markdown(response)