# Client python microservice for processing json documents before indexing

import json
import logging
import time

from flask import Flask, abort, jsonify, request

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


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[start:end]
        last_newline = chunk.rfind('\n')
        last_space = chunk.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):
        cleaned = {}
        for key, value in data.items():
            if not any(key.startswith(prefix) for prefix in prefixes):
                cleaned[key] = remove_fields_by_prefix(value, prefixes)
        return cleaned
    elif isinstance(data, list):
        return [remove_fields_by_prefix(item, prefixes) for item in data]
    else:
        return data


def process(client, data):
    logger.info('Processing client ' + client)

    if client == "WIKIPEDIA_001":
        try:
            title = data['title']
        except KeyError:
            abort(400, description="Missing 'title' field for WIKIPEDIA_001")
        url = "https://en.wikipedia.org/wiki/" + title
        data['url'] = url
        data['uri'] = client + "/" + url
        return jsonify(data)

    if client == "TREC_001":
        try:
            tabs = data['content'].split('\t')
        except KeyError:
            abort(400, description="Missing 'content' field for TREC_001")
        if len(tabs) < 2:
            return ""
        id = tabs[0]
        title = tabs[1]
        content = tabs[2] if len(tabs) > 2 else ""
        return jsonify({'id': id, 'title': title, 'content': content})

    if client == "TREC_002":
        try:
            id = data['id']
            contents = data['contents']
        except KeyError as e:
            abort(400, description=f"Missing field for TREC_002: {e}")
        doc = json.loads(contents)
        doc['id'] = id
        return jsonify(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),
                'title':        data.get('title', ''),
                'fileName':     data.get('fileName', ''),
                'Content-Type': data.get('Content-Type', ''),
                'content':      chunk,
            }
            docs.append(doc)
        return jsonify(docs)

    return jsonify(data)


app = Flask(__name__)
app.config['DEBUG'] = False


@app.route("/process", methods=['POST'])
def post_process():
    payload = request.get_json()
    if not payload or 'client' not in payload or 'data' not in payload:
        abort(400, description="Request must include 'client' and 'data' fields")
    start = time.perf_counter()
    output = process(payload['client'], payload['data'])
    elapsed_ms = int((time.perf_counter() - start) * 1000)
    logger.info('Processing took %d ms' % elapsed_ms)
    return output


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