
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#


import time
import threading
import queue
import psycopg2
from typing import Any, List, Dict, Tuple
from sanic import Sanic
from sanic.response import json
import calendar
import os
import logging

log_logger_folder_name = "log_cache_service"
if not os.path.exists(f"./{log_logger_folder_name}"):
    os.makedirs(f"./{log_logger_folder_name}")
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s %(levelname)-8s %(message)s',
                    datefmt='%d %b %Y %H:%M:%S',
                    filename=f"./{log_logger_folder_name}/log_{str(calendar.timegm(time.gmtime()))}", filemode='w')

USER = "postgres"
HOST = "127.0.0.1"
PORT = "28814"
DB_NAME = "pg_extension"
CACHE_SIZE = 10


class CacheService:
    def __init__(self, name_space: str, database: str, table: str, columns: List, batch_size: int, max_size: int = CACHE_SIZE):
        """
        name_space: train, valid, test
        database: database to use
        table: which table
        columns: selected cols
        max_size: max batches to cache
        """
        self.name_space = name_space
        self.batch_size = batch_size
        self.last_id = -1
        self.database = database
        self.table = table
        self.columns = columns
        self.queue = queue.Queue(maxsize=max_size)
        self.thread = threading.Thread(target=self.fetch_data, daemon=True)
        self.thread.start()

    def decode_libsvm(self, columns):
        map_func = lambda pair: (int(pair[0]), float(pair[1]))
        # 0 is id, 1 is label
        id, value = zip(*map(lambda col: map_func(col.split(':')), columns[2:]))
        sample = {'id': list(id),
                  'value': list(value),
                  'y': int(columns[1])}
        return sample

    def pre_processing(self, mini_batch_data: List[Tuple]):
        """
        mini_batch_data: [('0', '0', '123:123', '123:123', '123:123',)
        """
        sample_lines = len(mini_batch_data)
        feat_id = []
        feat_value = []
        y = []

        for i in range(sample_lines):
            row_value = mini_batch_data[i]
            sample = self.decode_libsvm(row_value)
            feat_id.append(sample['id'])
            feat_value.append(sample['value'])
            y.append(sample['y'])
        return {'id': feat_id, 'value': feat_value, 'y': y}

    def fetch_data(self):
        with psycopg2.connect(database=self.database, user=USER, host=HOST, port=PORT) as conn:
            while True:
                try:
                    # fetch and preprocess data from PostgreSQL
                    batch, time_usg = self.fetch_and_preprocess(conn)
                    self.queue.put(batch)
                    print(f"Data is fetched, {self.name_space} queue_size={self.queue.qsize()}, time_usg={time_usg}")
                    logger.info(f"Data is fetched, queue_size={self.queue.qsize()}, time_usg={time_usg}")
                    # block until a free slot is available
                    time.sleep(0.1)
                except psycopg2.OperationalError:
                    logger.exception("Lost connection to the database, trying to reconnect...")
                    time.sleep(5)  # wait before trying to establish a new connection
                    conn = psycopg2.connect(database=self.database, user=USER, host=HOST, port=PORT)

    def fetch_and_preprocess(self, conn):
        begin_time = time.time()
        cur = conn.cursor()
        # Assuming you want to get the latest 100 rows
        columns_str = ', '.join(self.columns)
        # Select rows greater than last_id
        cur.execute(f"SELECT id, {columns_str} FROM {self.table} "
                    f"WHERE id > {self.last_id} ORDER BY id ASC LIMIT {self.batch_size}")
        rows = cur.fetchall()

        if rows:
            # Update last_id with max id of fetched rows
            self.last_id = max(row[0] for row in rows)  # assuming 'id' is at index 0
        else:
            # If no more new rows, reset last_id to start over scan and return 'end_position'
            self.last_id = -1
            return "end_position", time.time() - begin_time

        batch = self.pre_processing(rows)
        return batch, time.time() - begin_time

    def get(self):
        return self.queue.get()

    def is_empty(self):
        return self.queue.empty()


app = Sanic("CacheServiceApp")


# start the server， this is from pg_ingerface
@app.route("/", methods=["POST"])
async def start_service(request):
    try:
        columns = request.json.get('columns')
        # can only be train or valid
        name_space = request.json.get('name_space')
        table_name = request.json.get('table_name')
        batch_size = request.json.get('batch_size')

        if columns is None:
            return json({"error": "No columns specified"}, status=400)
        if name_space not in ["train", "valid", "test"]:
            return json({"error": name_space + " is not correct"}, status=400)

        print(f"columns are {columns}, name_space = {name_space}")

        if not hasattr(app.ctx, f'{table_name}_{name_space}_cache'):
            setattr(app.ctx, f'{table_name}_{name_space}_cache',
                    CacheService(name_space, DB_NAME, table_name, columns, batch_size, CACHE_SIZE))

        return json("OK")
    except Exception as e:
        return json({"error": str(e)}, status=500)


# serve the data retrieve request from eva_service.py
@app.route("/", methods=["GET"])
async def serve_get_request(request):
    name_space = request.args.get('name_space')
    table_name = request.args.get('table_name')

    # check if exist
    if not hasattr(app.ctx, f'{table_name}_{name_space}_cache'):
        return json({"error": f"{table_name}_{name_space}_cache not start yet"}, status=404)

    # get data
    data = getattr(app.ctx, f'{table_name}_{name_space}_cache').get()

    # return
    if data is None:
        return json({"error": "No data available"}, status=404)
    else:
        return json(data)


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8093)