import streamlit as st # type: ignore
import os
from sentence_transformers import SentenceTransformer
from translate_app import tr
import getpass
from langchain_mistralai import ChatMistralAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import START, END, MessagesState, StateGraph
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from typing import Sequence
from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage, trim_messages
from langgraph.graph.message import add_messages
from typing_extensions import Annotated, TypedDict

import warnings
warnings.filterwarnings('ignore')

title = "Sales coaching"
sidebar_name = "Sales coaching"
dataPath = st.session_state.DataPath

os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_ENDPOINT"]="https://api.smith.langchain.com"
os.environ["LANGCHAIN_HUB_API_URL"]="https://api.smith.langchain.com"
os.environ["LANGCHAIN_PROJECT"] = "Sales Coaching Chatbot" 
os.getenv("LANGCHAIN_API_KEY")
os.getenv("MISTRAL_API_KEY")
model = ChatMistralAI(model="mistral-large-latest")

dataPath = st.session_state.DataPath

trimmer = trim_messages(
    max_tokens=1000,
    strategy="last",
    token_counter=model,
    include_system=True,
    allow_partial=False,
    start_on="human",
)

prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a helpful assistant. Answer all questions to the best of your ability in {language}.",
        ),
        MessagesPlaceholder(variable_name="messages"),
    ]
)

class State(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]
    language: str

def call_model(state: State):
    chain = prompt | model
    trimmed_messages = trimmer.invoke(state["messages"])
    response = chain.invoke(
        {"messages": trimmed_messages, "language": state["language"]}
    )
    return {"messages": [response]}

# Define a new graph
workflow = StateGraph(state_schema=State)

# Define the (single) node in the graph
workflow.add_edge(START, "model")
workflow.add_node("model", call_model)
workflow.add_edge("model", END)

# Add memory
memory = MemorySaver()
app = workflow.compile(checkpointer=memory)

config = {"configurable": {"thread_id": "abc123"}}

def run():

    st.write("")
    st.write("")
    st.title(tr(title))
    
    context = """Simuler conversation entre un commercial de mon entreprise, Marc, et interlocuteur chez un de mes prospect. 
Chez mon prospect, tu es le Directeur commercial.  
Tu simules les réponses du Directeur commercial.  
Moi, je saisie les questions du Commercial de mon entreprise, Marc.
Mon entreprise propose une solution logicielle pour gérer la proposition de valeur d’entreprises B2B qui commercialises des solutions technologiques.   
  
Le prospect est dans (Profil des clients cibles) suivant:  
 ·      Les scale-up françaises qui commercialisent une solution technologique B2B  
 ·      Les sociétés françaises qui commercialisent une solution technologique B2B   
   

  
Problématique adressée par ma (celle de Marc) solution:  
 ·      Il est difficile pour les startups de concevoir et formaliser une proposition de valeur unique et pertinente
 ·      Il est difficile pour les forces de vente d'articuler clairement les messages de la proposition de valeur
 ·      Il est chronophage pour les scale-ups de former leur forces de ventes sur les évolutions de la proposition de valeur et de ses messages   Valeur ajoutée de ma solution  
 ·      augmenter les performances commerciales
 ·      Réduction du cycle de vente
 ·      Augmenter la croissance du chiffre d'affaires
 ·      Concentrer les ressources sur les opportunités qualifiées
 ·      Réduction temps de monté en compétence des nouvelles embauches
 ·      augmenter le taux de conversion d'affaires gagnées
 ·      améliorer l'efficacité globale des ventes
 ·      Améliorer l'efficience et la confiance des forces de ventes  
 ·      optimiser les processus commerciaux pour réduire les cycles de vente, fidéliser les clients et augmenter la productivité
 ·      favorise le développement personnel des forces de vente   

Les cas d’usages adressés par ma solution  
 ·      Affiner et modéliser la proposition de valeur
 ·      Décliner les messages marketing & commerciaux
 ·      Former la force de vente sur la proposition de valeur
 ·      Orchestrer les conversations clients
 ·      Partager les retours terrain
 ·      Mettre en oeuvre des meilleures pratiques commerciales
 ·      Identifier et reproduire les messages de vente gagnants
"""
  
    messages = [
        SystemMessage(content=context)
        ]
    
    trimmer.invoke(messages)
    
    # query = "Hi I'm Todd, please tell me a joke."
    query = st.text_area(label=tr("Vendeur:"), value="")
    language = "French"
    st.button(label=tr("Validez"), type="primary")
    
    input_messages = [HumanMessage(query)]
    output = app.invoke(
            {"messages": input_messages, "language": language},
            config,
        )
    st.write(output["messages"][-1].content) #.pretty_print())

    
    '''
    # Créer un espace réservé pour afficher les tokens
    placeholder = st.empty()
    
    for chunk, metadata in app.stream(
        {"messages": input_messages, "language": language},
        config,
        stream_mode="messages",
        ):
        if isinstance(chunk, AIMessage):  # Filter to just model responses
            # st.markdown("<span style='white-space: nowrap;'>"+"/"+chunk.content+"/"+"</span>", unsafe_allow_html=True)
            placeholder.markdown(f"<p style='display: inline;'>{chunk.content}</p>", unsafe_allow_html=True)
    '''
    st.write("")
    st.write("")
    st.write("")
    st.write("")