from os import makedirs, remove
from os.path import exists, dirname
from functools import cache
import json
import streamlit as st
from googleapiclient.discovery import build
from slugify import slugify
from transformers import pipeline
import uuid

from beautiful_soup.beautiful_soup import get_url_content

"""
Request Google Search API with query and return results.
"""
@cache
def google_search_api_request( query ):
    api_key = st.secrets["google_search_api_key"]
    cx = st.secrets["google_search_engine_id"]
    service = build(
        "customsearch",
        "v1",
        developerKey=api_key,
        cache_discovery=False
    )

    # Exclude PDFs from search results.
    query = query + ' -filetype:pdf'

    return service.cse().list(
        q=query,
        cx=cx,
        num=5,
        lr='lang_en', # lang_de
        fields='items(title,link),searchInformation(totalResults)'
        ).execute()

"""
Request Google Search API with query and return results. Results are cached in files.
"""
def search_results( query ):
    file_path = 'search-results/' + slugify( query ) + '.json'

    results = []
    makedirs(dirname(file_path), exist_ok=True)
    if exists( file_path ):
        with open( file_path, 'r' ) as results_file:
            results = json.load( results_file )
    else:
        search_result = google_search_api_request( query )
        if int( search_result['searchInformation']['totalResults'] ) > 0:
            results = search_result['items']
            with open( file_path, 'w' ) as results_file:
                json.dump( results, results_file )

    if len( results ) == 0:
        raise Exception('No results found.')
    
    return results

"""
Generate summary for content.
"""
def generate_summary( url_id, content ):
    file_path = 'summaries/' + url_id + '.json'
    makedirs(dirname(file_path), exist_ok=True)
    if exists( file_path ):
        with open( file_path, 'r' ) as file:
            summary = json.load( file )
    else:
        try:
            summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
            # https://huggingface.co/docs/transformers/v4.18.0/en/main_classes/pipelines#transformers.SummarizationPipeline
            summary = summarizer(content, max_length=130, min_length=30, do_sample=False, truncation=True)
        except Exception as exception:
            raise exception

        with open( file_path, 'w' ) as file:
            json.dump( summary, file )
    
    return summary

"""
Helper function for exception notices.
"""
def exception_notice( exception ):
    query_params = st.experimental_get_query_params()
    if 'debug' in query_params.keys() and query_params['debug'][0] == 'true':
        st.exception(exception)
    else:
        st.warning(str(exception))

"""
Checks if string contains keyword.
"""
def is_keyword_in_string( keywords, string ):
    for keyword in keywords:
        if keyword in string:
            return True
    return False

def filter_strings_by_keywords( strings, keywords ):
    content = ''
    for string in strings:
        # Filter strings with keywords
        if is_keyword_in_string( keywords, string ):
            content += string + '\n'
    return content

def main():
    st.title('Racoon Search')
    query = st.text_input('Search query')
    query_params = st.experimental_get_query_params()

    if query :
        with st.spinner('Loading search results...'):
            try:
                results = search_results( query )
            except Exception as exception:
                exception_notice(exception)
                return

        number_of_results = len( results )
        st.success( 'Found {} results for "{}".'.format( number_of_results, query ) )

        if 'debug' in query_params.keys() and query_params['debug'][0] == 'true':
            with st.expander("Search results JSON"):
                if st.button('Delete search result cache', key=query + 'cache'):
                    remove( 'search-results/' + slugify( query ) + '.json' )
                st.json( results )

        progress_bar = st.progress(0)

        st.header('Search results')
        st.markdown('---')

        # for result in results:
        for index, result in enumerate(results):
            with st.container():
                st.markdown('### ' + result['title'])
                url_id = uuid.uuid5( uuid.NAMESPACE_URL, result['link'] ).hex
                try:
                    strings = get_url_content( result['link'] )
                    keywords = query.split(' ')
                    content = filter_strings_by_keywords( strings, keywords )
                    # print(content)
                    # print(len(content.split()))
                    summary = generate_summary( url_id, content )
                    for sentence in summary:
                        st.write(sentence['summary_text'])
                except Exception as exception:
                    exception_notice(exception)

                progress_bar.progress( ( index + 1 ) / number_of_results )

                col1, col2, col3 = st.columns(3)
                with col1:
                    st.markdown('[Website Link]({})'.format(result['link']))

                with col2:
                    if st.button('Delete content from cache', key=url_id + 'content'):
                        remove( 'page-content/' + url_id + '.txt' )

                with col3:
                    if st.button('Delete summary from cache', key=url_id + 'summary'):
                        remove( 'summaries/' + url_id + '.json' )

                st.markdown('---')
 

if __name__ == '__main__':
    main()