# Client python microservice for processing json documents before indexing
#
import os
import logging
import time
import json
from flask import Flask, abort, jsonify, request

def chunk_text(text, chunk_size, overlap=0):
    chunks = []
    start = 0
    text_len = len(text)

    while start < text_len:
        end = start + chunk_size
        if end >= text_len:
            chunks.append(text[start:])
            break

        chunk_text = text[start:end]
        last_newline = chunk_text.rfind('\n')
        last_space = chunk_text.rfind(' ')
        break_pos = max(last_newline, last_space)

        if break_pos != -1:
            actual_end = start + break_pos + 1
            chunks.append(text[start:actual_end])

            if overlap > 0:
                overlap_start = actual_end - overlap
                if overlap_start > start:
                    while overlap_start < actual_end and text[overlap_start] not in (' ', '\n'):
                        overlap_start += 1
                    if overlap_start < actual_end:
                        overlap_start += 1
                    start = overlap_start
                else:
                    start = actual_end
            else:
                start = actual_end
        else:
            chunks.append(text[start:end])
            start = end - overlap

    return chunks


def remove_fields_by_prefix(data, prefixes):
     if isinstance(data, dict):
         # Create a new dict with only fields that don't match prefixes
         cleaned = {}
         for key, value in data.items():
             # Check if key starts with any of the prefixes
             if not any(key.startswith(prefix) for prefix in prefixes):
                 # Recursively clean the value
                 cleaned[key] = remove_fields_by_prefix(value, prefixes)
         return cleaned

     elif isinstance(data, list):
         # Recursively clean each item in the list
         return [remove_fields_by_prefix(item, prefixes) for item in data]

     else:
         # Primitive value, return as-is
         return data

def process(client, data):
    logger.info('Processing client ' + client)
    
    if (client == "WIKIPEDIA_001"):
        title = data['title']
        url = "https://en.wikipedia.org/wiki/" + title
        data['url'] = url
        uri = client + "/" + url
        data['uri'] = uri

#        content = data['content']
#        content = content[0:1000] # cut to 1000 length 
#        data['content'] = content
   
        return jsonify(data)

    if (client == "TREC_001"):
        tabs = data['content'].split('\t')
        if len(tabs) < 2:
            return ""
        id = tabs[0]
        title = tabs[1]
        if len(tabs) > 2:
            content = tabs[2]
        else :
            content = ""
        return jsonify({'id': id, 'title': title, 'content': content})

    if (client == "TREC_002"):
        id = data['id']
        contents = data['contents']
        doc = json.loads(contents)
        doc['id'] = id
#        doc = remove_fields_by_prefix(doc, ["info_", "attrs_"])
        return doc

    if (client == "EPS_001"):
        docs = []
        chunks = chunk_text(data.get('content', ''), 10000, overlap=100)
        for idx, chunk in enumerate(chunks):
            doc = {'id': data.get('fileName', '') + '#' + str(idx)}
            doc['title'] = data.get('title', '')
            doc['fileName'] = data.get('fileName', '')
            doc['Content-Type'] = data.get('Content-Type', '')
            doc['content'] = chunk
            docs.append(doc)

        return docs

    return jsonify(data)


logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO,
)
logger = logging.getLogger(__name__)

app = Flask(__name__)
app.config['CORS_HEADERS'] = 'Content-Type'
app.config['DEBUG'] = False

@app.route("/process", methods=['POST'])
def post_process():
    input = request.get_json()
    if 'client' not in input or 'data' not in input:
        abort(500)
    else:
        start = time.time()
        output = process(input['client'], input['data'])
        end = time.time() 
        logger.info('Processing took %d (ms)' % int((end - start)*1000))
        return output

if __name__ == "__main__":
    app.run(host='0.0.0.0')

