import streamlit as st
import spacy
import wikipediaapi
import wikipedia
from wikipedia.exceptions import DisambiguationError
from transformers import TFAutoModel, AutoTokenizer
import numpy as np
import pandas as pd
import faiss
import datetime
import time


try:
    nlp = spacy.load("en_core_web_sm")
except:
    spacy.cli.download("en_core_web_sm")
    nlp = spacy.load("en_core_web_sm")

wh_words = ['what', 'who', 'how', 'when', 'which']

def get_concepts(text):
    text = text.lower()
    doc = nlp(text)
    concepts = []
    for chunk in doc.noun_chunks:
        if chunk.text not in wh_words:
            concepts.append(chunk.text)
    return concepts

def get_passages(text, k=100):
    doc = nlp(text)
    passages = []
    passage_len = 0
    passage = ""
    sents = list(doc.sents)
    for i in range(len(sents)):
        sen = sents[i]
        passage_len += len(sen)
        if passage_len >= k:
            passages.append(passage)
            passage = sen.text
            passage_len = len(sen)
            continue
        elif i == (len(sents) - 1):
            passage += " " + sen.text
            passages.append(passage)
            passage = ""
            passage_len = 0
            continue
        passage += " " + sen.text
    return passages

def get_dicts_for_dpr(concepts, n_results=20, k=100):
    dicts = []
    for concept in concepts:
        wikis = wikipedia.search(concept, results=n_results)
        st.write(f"{concept} No of Wikis: {len(wikis)}")
        for wiki in wikis:
            try:
                html_page = wikipedia.page(title=wiki, auto_suggest=False)
            except DisambiguationError:
                continue
            htmlResults = html_page.content
            passages = get_passages(htmlResults, k=k)
            for passage in passages:
                i_dicts = {}
                i_dicts['text'] = passage
                i_dicts['title'] = wiki
                dicts.append(i_dicts)
    return dicts

passage_encoder = TFAutoModel.from_pretrained("nlpconnect/dpr-ctx_encoder_bert_uncased_L-2_H-128_A-2")
query_encoder = TFAutoModel.from_pretrained("nlpconnect/dpr-question_encoder_bert_uncased_L-2_H-128_A-2")
p_tokenizer = AutoTokenizer.from_pretrained("nlpconnect/dpr-ctx_encoder_bert_uncased_L-2_H-128_A-2")
q_tokenizer = AutoTokenizer.from_pretrained("nlpconnect/dpr-question_encoder_bert_uncased_L-2_H-128_A-2")

def get_title_text_combined(passage_dicts):
    res = []
    for p in passage_dicts:
        res.append(tuple((p['title'], p['text'])))
    return res

def extracted_passage_embeddings(processed_passages, max_length=156):
    passage_inputs = p_tokenizer.batch_encode_plus(
                    processed_passages,
                    add_special_tokens=True,
                    truncation=True,
                    padding="max_length",
                    max_length=max_length,
                    return_token_type_ids=True
                )
    passage_embeddings = passage_encoder.predict([np.array(passage_inputs['input_ids']), np.array(passage_inputs['attention_mask']), 
                                            np.array(passage_inputs['token_type_ids'])], 
                                            batch_size=64, 
                                            verbose=1)
    return passage_embeddings

def extracted_query_embeddings(queries, max_length=64):
    query_inputs = q_tokenizer.batch_encode_plus(
        queries,
        add_special_tokens=True,
        truncation=True,
        padding="max_length",
        max_length=max_length,
        return_token_type_ids=True
    )
    
    query_embeddings = query_encoder.predict([np.array(query_inputs['input_ids']),
        np.array(query_inputs['attention_mask']),
        np.array(query_inputs['token_type_ids'])],
        batch_size=1,
        verbose=1)
    return query_embeddings

def get_pagetext(page):
    s = str(page).replace("/t","")
    return s

def get_wiki_summary(search):
    wiki_wiki = wikipediaapi.Wikipedia('en')
    page = wiki_wiki.page(search)                                   


def get_wiki_summaryDF(search):
    wiki_wiki = wikipediaapi.Wikipedia('en')
    page = wiki_wiki.page(search)

    isExist = page.exists()
    if not isExist:
        return isExist, "Not found", "Not found", "Not found", "Not found"

    pageurl = page.fullurl
    pagetitle = page.title
    pagesummary = page.summary[0:60]
    pagetext = get_pagetext(page.text)

    backlinks = page.backlinks
    linklist = ""
    for link in backlinks.items():
      pui = link[0]
      linklist += pui + " ,  "
      a=1 
      
    categories = page.categories
    categorylist = ""
    for category in categories.items():
      pui = category[0]
      categorylist += pui + " ,  "
      a=1     
    
    links = page.links
    linklist2 = ""
    for link in links.items():
      pui = link[0]
      linklist2 += pui + " ,  "
      a=1 
      
    sections = page.sections
    
    ex_dic = {
      'Entity' : ["URL","Title","Summary", "Text", "Backlinks", "Links", "Categories"],
      'Value': [pageurl, pagetitle, pagesummary, pagetext, linklist,linklist2, categorylist ]
    }

    df = pd.DataFrame(ex_dic)
    
    return df


def save_message(name, message):
    now = datetime.datetime.now()
    timestamp = now.strftime("%Y-%m-%d %H:%M:%S")
    with open("chat.txt", "a") as f:
        f.write(f"{timestamp} - {name}: {message}\n")

def main():
    st.title("Streamlit Chat")

    name = st.text_input("Name")
    message = st.text_input("Message")
    if st.button("Submit"):
        
        # wiki
        df = get_wiki_summaryDF(message)
        
        save_message(name, message)
        save_message(name, df)
        
        st.text("Message sent!")

    
    st.text("Chat history:")
    with open("chat.txt", "a+") as f:
        f.seek(0)
        chat_history = f.read()
    st.text(chat_history)

    countdown = st.empty()
    t = 60
    while t:
        mins, secs = divmod(t, 60)
        countdown.text(f"Time remaining: {mins:02d}:{secs:02d}")
        time.sleep(1)
        t -= 1
        if t == 0:
            countdown.text("Time's up!")
            with open("chat.txt", "a+") as f:
                f.seek(0)
                chat_history = f.read()
            st.text(chat_history)
            t = 60

if __name__ == "__main__":
    main()