from langchain_openai import ChatOpenAI#, OpenAIEmbeddings # No need to pay for using embeddings as well when have free alternatives

# Data
from langchain_community.document_loaders import DirectoryLoader, TextLoader, WebBaseLoader
# from langchain_chroma import Chroma # The documentation uses this one, but it is extremely recent, and the same functionality is available in langchain_community and langchain (which imports community)
from langchain_community.vectorstores import Chroma # This has documentation on-hover, while the indirect import through non-community does not
from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings # The free alternative (also the default in docs, with model_name = 'all-MiniLM-L6-v2')
from langchain.text_splitter import RecursiveCharacterTextSplitter#, TextSplitter # Recursive to better keep related bits contiguous (also recommended in docs: https://python.langchain.com/docs/modules/data_connection/document_transformers/)

# Chains
from langchain.prompts import PromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain.tools.retriever import create_retriever_tool
from langchain_core.runnables import RunnablePassthrough, RunnableParallel

# Agents
from langchain import hub
from langchain.agents import create_tool_calling_agent, AgentExecutor

# To manually create inputs to test pipelines
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.documents import Document

# To serve the app
from fastapi import FastAPI
from langchain.pydantic_v1 import BaseModel, Field
from langchain_core.messages import BaseMessage
from langserve import add_routes, CustomUserType

import requests
from bs4 import BeautifulSoup

import os
import shutil
import re

import dotenv
dotenv.load_dotenv()


## Vector stores

script_db = Chroma(embedding_function = SentenceTransformerEmbeddings(model_name = 'all-MiniLM-L6-v2'), persist_directory = r'scripts\db')
woo_db = Chroma(embedding_function = SentenceTransformerEmbeddings(model_name = 'all-MiniLM-L6-v2'), persist_directory = r'wookieepedia_db')



## Wookieepedia functions

def first_wookieepedia_result(query: str) -> str:
    '''Get the url of the first result when searching Wookieepedia for a query
    (best for simple names as queries, ideally generated by the llm for something like
    "Produce a input consisting of the name of the most important element in the query so that its article can be looked up")
    '''
    search_results = requests.get(f'https://starwars.fandom.com/wiki/Special:Search?query={"+".join(query.split(" "))}')
    soup = BeautifulSoup(search_results.content, 'html.parser')
    first_res = soup.find('a', class_ = 'unified-search__result__link')
    return first_res['href']


def get_new_wookieepedia_chunks(query: str, previous_sources: set[str]) -> list[Document]:
    '''Retrieve and return chunks of the content of the first result of query on Wookieepedia, then return the closest matches for.
    '''
    url = first_wookieepedia_result(query)

    if url in previous_sources: return []
    else:
        doc = WebBaseLoader(url).load()[0] # Only one url passed in => only one Document out; no need to assert
    
        # There probably is a very long preamble before the real content, however, if more than one gap then ignore and proceed with full document
        trimmed = parts[1] if len(parts := doc.page_content.split('\n\n\n\n\n\n\n\n\n\n\n\n\n\n \xa0 \xa0')) == 2 else doc.page_content
        doc.page_content = re.sub(r'[\n\t]{2,}', '\n', trimmed) # And remove excessive spacing

        return RecursiveCharacterTextSplitter(chunk_size = 800, chunk_overlap = 100).split_documents([doc])

def get_wookieepedia_context(original_query: str, simple_query: str, wdb: Chroma) -> list[Document]:
    try:
        new_chunks = get_new_wookieepedia_chunks(simple_query, previous_sources = set(md.get('source') for md in wdb.get()['metadatas']))
        if new_chunks: wdb.add_documents(new_chunks)
    except: return []

    return wdb.similarity_search(original_query, k = 10)



# Chains

llm = ChatOpenAI(model = 'gpt-3.5-turbo-0125', temperature = 0)


document_prompt_system_text = '''
You are very knowledgeable about Star Wars and your job is to answer questions about its plot, characters, etc.
Use the context below to produce your answers with as much detail as possible.
If you do not know an answer, say so; do not make up information not in the context.

<context>
{context}
</context>
'''
document_prompt = ChatPromptTemplate.from_messages([
    ('system', document_prompt_system_text),
    MessagesPlaceholder(variable_name = 'chat_history', optional = True),
    ('user', '{input}')
])
document_chain = create_stuff_documents_chain(llm, document_prompt)


script_retriever_prompt = ChatPromptTemplate.from_messages([
    MessagesPlaceholder(variable_name = 'chat_history'),
    ('user', '{input}'),
    ('user', '''Given the above conversation, generate a search query to look up relevant information in a database containing the full scripts from the Star Wars films (i.e. just dialogue and brief scene descriptions).
     The query need not be a proper sentence, but a list of keywords likely to be in dialogue or scene descriptions''')
])
script_retriever_chain = create_history_aware_retriever(llm, script_db.as_retriever(), script_retriever_prompt) # Essentially just: prompt | llm | StrOutputParser() | retriever


woo_retriever_prompt = ChatPromptTemplate.from_messages([
    MessagesPlaceholder(variable_name = 'chat_history'),
    ('user', '{input}'),
    ('user', 'Given the above conversation, generate a search query to find a relevant page in the Star Wars fandom wiki; the query should be something simple, such as the name of a character, place, event, item, etc.')
])
woo_retriever_chain = create_history_aware_retriever(llm, woo_db.as_retriever(), woo_retriever_prompt) # Essentially just: prompt | llm | StrOutputParser() | retriever


full_chain = create_retrieval_chain(script_retriever_chain, document_chain)



## Agent version

script_tool = create_retriever_tool(
    script_db.as_retriever(search_kwargs = dict(k = 4)),
    'search_film_scripts',
    '''Search the Star Wars film scripts. This tool should be the first choice for Star Wars related questions.
    Queries passed to this tool should be lists of keywords likely to be in dialogue or scene descriptions, and should not include film titles.'''
)
wookieepedia_tool = create_retriever_tool(
    woo_db.as_retriever(search_kwargs = dict(k = 4)),
    'search_wookieepedia',
    'Search the Star Wars fandom wiki. This tool should be used for queries about details of a particular character, location, event, weapon, etc., and the query should be something simple, such as the name of a character, place, event, item, etc.',
)
tools = [script_tool, wookieepedia_tool]

agent_system_text = '''
You are a helpful agent who is very knowledgeable about Star Wars and your job is to answer questions about its plot, characters, etc.
Use the context provided in the exchanges to come to produce your answers with as much detail as possible.
If you do not know an answer, say so; do not make up information.
'''
agent_prompt = ChatPromptTemplate.from_messages([
    ('system', agent_system_text),
    MessagesPlaceholder('chat_history', optional = True), # Using this form since not clear how to have optional = True in the tuple form
    ('human', '{input}'),
    ('placeholder', '{agent_scratchpad}') # Required for chat history and the agent's intermediate processing values
])
agent = create_tool_calling_agent(llm, tools, agent_prompt)
agent_executor = AgentExecutor(agent = agent, tools = tools, verbose = True)




class StrInput(BaseModel):
    input: str

class Input(BaseModel):
    input: str
    chat_history: list[BaseMessage] = Field(
        ...,
        extra = {'widget': {'type': 'chat', 'input': 'location'}},
    )

class Output(BaseModel):
    output: str




## App definition

app = FastAPI(
  title = 'LangChain Server',
  version = '1.0',
  description = 'Simple version of a Star Wars expert',
)



## Adding chain route

# add_routes(app, script_db.as_retriever())
add_routes(app, full_chain.with_types(input_type = StrInput, output_type = Output), playground_type = 'default')

# NOTE: The chat playground type has a web page issue (flashes and becomes white, hence non-interactable; this was supposedly solved in an issue late last year)

# add_routes(app, agent_executor, playground_type = 'chat')
# add_routes(app, agent_executor.with_types(input_type = StrInput, output_type = Output))


if __name__ == '__main__':
    import uvicorn

    uvicorn.run(app, host = 'localhost', port = 8000)