from concurrent import futures
from concurrent.futures import ProcessPoolExecutor

def download_parallel_multiprocessing():
   with ProcessPoolExecutor() as executor:
       future_to_key = {executor.submit(download_one_file, key): key for key in list_key_to_download}

       for future in futures.as_completed(future_to_key):
           key = future_to_key[future]
           exception = future.exception()

           if not exception:
               yield key, future.result()
           else:
               yield key, exception

for key, result in download_parallel_multiprocessing():
#     print(f"{key}: {result}")
   pass

from langchain_community.llms import WatsonxLLM
# from dotenv import load_dotenv
import os
from ibm_watson_machine_learning.metanames import GenTextParamsMetaNames as GenParams

# Model Declaration

api_key = "Your Key Here"
ibm_cloud_url = "Write your URL"
project_id = "The project id"

# Params Declaration
params = {
   GenParams.DECODING_METHOD: "sample",
   GenParams.MIN_NEW_TOKENS: 50,
   GenParams.MAX_NEW_TOKENS: 430,
   GenParams.RANDOM_SEED: 42,
   GenParams.TEMPERATURE: 0,
   GenParams.TOP_K: 20,
   GenParams.TOP_P:1
}
# Model Configuration
llm = WatsonxLLM(
           model_id='intfloat/multilingual-e5-large',
           url=ibm_cloud_url,
           apikey=api_key,
           project_id=project_id,
           params=params,
           )

import re

def matching_files_with_names(file_names):
   matching_files = []
   for file_name in file_names:
       # Check if the file name contains a pattern
       pattern_match = re.search(r'\[\d+\]', file_name)
       if pattern_match:
           # Remove the pattern and check if the resulting name exists in the list
           pattern_removed_name = re.sub(r'\[\d+\].', '', file_name)
           if pattern_removed_name in file_names and pattern_removed_name != file_name:
               matching_files.append([file_name, pattern_removed_name])
   return matching_files

matching_files = matching_files_with_names(matching_pairs)

import lxml
import pickle
from lxml import etree
from langchain.docstore.document import Document

dir_list_pre = os.listdir("temp_pre/")
# destination_path = "temp_post/"
full_doc = []

for file in tqdm(list(matching_pairs_og.keys())):
   file_path = "temp_pre/" + file
   with open(file_path, "r") as f:
       xml_content = f.read()

   # Parse the XML
   try:
       with open(file_path, "rb") as f:
           xml_content_tag_reader = f.read()
       parser = etree.XMLParser(recover=True)
       root = etree.fromstring(xml_content_tag_reader, parser)
       # Get the ElementTree object
       tree = root.getroottree()
       #Extracting metadata except that of text
       meta_tags = {tree.getpath(d): d.text for d in root.iterdescendants() if '/txt' not in tree.getpath(d)}

       soup = BeautifulSoup(xml_content, "xml")

       text_tags = soup.find_all('txt')
       html_content = ''.join([' ' + tag.get_text() for tag in text_tags])

       if len(html_content)>0:

           html_soup = BeautifulSoup(html_content, "html.parser")

           # Find all <p> tags containing <strong> tags
           paragraphs_with_strong = html_soup.find_all(lambda tag: tag.name == 'p' and tag.strong)

           # Iterate through each <p> tag containing <strong> tags
           for p_tag in paragraphs_with_strong:
               strong_tag = p_tag.strong
               if strong_tag:
                   strong_text = strong_tag.text.strip()
                   # Check if the text starts with a digit followed by a period and a space
                   if re.match(r'^\d+\. ', strong_text):
                       # Find the previous sibling header tag
                       previous_header = p_tag.find_previous_sibling(lambda tag: tag.name.startswith('h') and len(tag.name) == 2)
                       if previous_header:
                           # Determine the level of the next header
                           next_header_level = 'h6'
                           # Replace the <p> tag with the next level header
                           new_header_tag = html_soup.new_tag(next_header_level)
                           new_header_tag.string = strong_text  # Retain the text content within <strong>
                           p_tag.replace_with(new_header_tag)

           # Get the modified HTML
           modified_html = str(html_soup)

           headers_to_split_on = [
               ("h1", "Header 1"),
               ("h2", "Header 2"),
               ("h3", "Header 3"),
               ("h4", "Header 4"),
               ("h5", "Header 5"),
               ("h6", "Header 6")
           ]

           html_splitter = HTMLHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
           html_header_splits = html_splitter.split_text(modified_html)
           for i, split in enumerate(html_header_splits):
               tokens = llm.get_num_tokens(split.page_content)
#                 html_header_splits[i].metadata.update(meta_tags)
               html_header_splits[i].metadata["token_count"] = tokens
               html_header_splits[i].metadata["source"] = file
               html_header_splits[i].metadata["index"] = i

           full_doc.append(html_header_splits)

       else:
           
           # Getting text from the file
           soup = BeautifulSoup(xml_content, "xml")
           text_tags = soup.find_all('txtascii')
           html_content = ''.join([' ' + tag.get_text() for tag in text_tags])
               
           lc_doc = Document(page_content=html_content)
           tokens = llm.get_num_tokens(html_content)
#             lc_doc.metadata["tags"] = meta_tags
           lc_doc.metadata["token_count"] = tokens
           lc_doc.metadata["source"] = file
           lc_doc.metadata["index"] = 1
           full_doc.append(lc_doc)
           
   except lxml.etree.XSLTApplyError as e:
       continue
   
# Save the list to a file
with open('full_xml.pkl', 'wb') as file:
   pickle.dump(full_doc, file)
   
#Upload the file to COS
cos_client.upload_file('./full_xml.pkl', bucket_name, 'full_xml.pkl')
# cos_client.Object(bucket_name, item_name).put(Body=filelike_object)

with open(file_path, "rb") as f:
           xml_content_tag_reader = f.read()
       parser = etree.XMLParser(recover=True)
       root = etree.fromstring(xml_content_tag_reader, parser)
       # Get the ElementTree object
       tree = root.getroottree()
       #Extracting metadata except that of text
       meta_tags = {tree.getpath(d): d.text for d in root.iterdescendants() if '/txt' not in tree.getpath(d)}

from langchain_text_splitters import HTMLHeaderTextSplitter
headers_to_split_on = [
   ("h1", "Header 1"),
   ("h2", "Header 2"),
   ("h3", "Header 3"),
   ("h4", "Header 4"),
   ("h5", "Header 5"),
   ("h6", "Header 6")
]

html_splitter = HTMLHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
html_header_splits = html_splitter.split_text(modified_html)

# Find all <p> tags containing <strong> tags
paragraphs_with_strong = html_soup.find_all(lambda tag: tag.name == 'p' and tag.strong)

# Iterate through each <p> tag containing <strong> tags
for p_tag in paragraphs_with_strong:
   strong_tag = p_tag.strong
   if strong_tag:
       strong_text = strong_tag.text.strip()
       # Check if the text starts with a digit followed by a period and a space
       if re.match(r'^\d+\. ', strong_text):
           # Find the previous sibling header tag
           previous_header = p_tag.find_previous_sibling(lambda tag: tag.name.startswith('h') and len(tag.name) == 2)
           if previous_header:
               # Determine the level of the next header
               next_header_level = 'h6'
               # Replace the <p> tag with the next level header
               new_header_tag = html_soup.new_tag(next_header_level)
               new_header_tag.string = strong_text  # Retain the text content within <strong>
               p_tag.replace_with(new_header_tag)

def extract_headers_and_content(html):
   html_soup = BeautifulSoup(html, 'html.parser')
   headers = ['h1', 'h2', 'h3', 'h4', 'h5', 'h6']
   result = []
   current_metadata = {}

   for tag in html_soup.find_all(headers):
       # Update the current metadata hierarchy based on the current tag
       level = headers.index(tag.name)
       current_metadata = {h: current_metadata[h] for h in headers[:level] if h in current_metadata}
       current_metadata[tag.name] = tag.get_text(strip=True)
       
       # Collect the content under the current header
       content = tag.find_next_sibling()
       content_text = ""
       while content and content.name not in headers:
           content_text += content.get_text(strip=True) + " "
           content = content.find_next_sibling()

       if content_text.strip():
           result.append({'page_content': content_text.strip(), 'metadata': current_metadata.copy()})

   return result

exception_file_chunking = []
for key, value in tqdm(full_doc_v2.items()):
   extracted_data = extract_headers_and_content(value)
   extracted_data_v2 = [i for i in extracted_data if i['page_content'] != '']
   for idx, i in enumerate(extracted_data_v2):
       i['metadata']['source'] = key
       tokens = llm.get_num_tokens(i['page_content'])
       i['metadata']['token_count'] = tokens
       i['metadata']['index'] = idx
   exception_file_chunking.append(extracted_data_v2)

#Renaming the keys in the dictionary

for k, v in tqdm(tags.items()):
   new_dict = {}
   for key, value in v.items():
       if key == '/Segmente/Segment':
           continue
       parts = key.rsplit('/', 2)
       text = '_'.join([part.lower() for part in parts[-2:]])
       new_dict[text] = value
   tags[k] = new_dict

from langchain.embeddings import HuggingFaceEmbeddings

model_name = "intfloat/multilingual-e5-large"
model_kwargs = {'device': 'cuda'}
encode_kwargs = {'normalize_embeddings': False}
hf = HuggingFaceEmbeddings(
   model_name=model_name,
   model_kwargs=model_kwargs,
   encode_kwargs=encode_kwargs
)

from elasticsearch import Elasticsearch
from langchain_elasticsearch import ElasticsearchStore

# Create a custom Elasticsearch client
es_client = Elasticsearch(
   es_url,
   basic_auth=(es_user, es_password),
   verify_certs=False,
   request_timeout = 1200
)

elastic_vector_search = ElasticsearchStore(
   index_name="index_e5_multiprocessing_htmlsplitter_gpu_test_v2",
   es_connection=es_client,
   embedding=hf,
)

db = ElasticsearchStore.from_documents(
  documents=texts,
   embedding=hf,
   index_name="index_e5_multiprocessing-v6",
   strategy=ElasticsearchStore.ApproxRetrievalStrategy(
       hybrid = True
   ),
   es_connection=es_client,
)