# /*
#  * Copyright 2024 KylinSoft Co., Ltd.
#  *
#  * This program is free software: you can redistribute it and/or modify it under
#  * the terms of the GNU General Public License as published by the Free Software
#  * Foundation, either version 3 of the License, or (at your option) any later
#  * version.
#  *
#  * This program is distributed in the hope that it will be useful, but WITHOUT
#  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
#  * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
#  * details.
#  *
#  * You should have received a copy of the GNU General Public License along with
#  * this program. If not, see <https://www.gnu.org/licenses/>.
#  */

from concurrent import futures
import grpc
import document_pb2_grpc
import document_pb2
from async_cache import AsyncCache
import logging
import json
import os
import multiprocessing
from logging.handlers import TimedRotatingFileHandler
from datetime import datetime, timedelta

logger = logging.getLogger('parse_service_logger')
logger.setLevel(logging.INFO)
formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")

console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)
console_handler.setFormatter(formatter)
logger.addHandler(console_handler)

logger.info("Logging is set up. Starting the service...")

class DocumentService(document_pb2_grpc.DocumentServiceServicer):
    def __init__(self, max_workers=4):
        self.max_workers = max_workers
        self.async_cache = AsyncCache()
        
    @staticmethod
    def parse_and_chunk_file(file_path, chunker_type, chunk_size,callbake_data):
        from parser_factory import ParserFactory
        from chunker_factory import ChunkerFactory
        """文件解析和分块逻辑，供进程池执行。"""
        try:
            parser = ParserFactory.create_parser(os.path.splitext(file_path)[1][1:])
            parsed_content = parser.parser(file_path)
            chunker = ChunkerFactory.create_chunker(chunker_type)
            chunks = chunker.segment_tokenize(parsed_content, chunk_size)
            return (file_path, chunks,callbake_data)
        except Exception as e:
            logger.error(f"Error parsing or chunking {file_path}: {str(e)}")
            return (file_path, [],callbake_data)

    def SyncChunkDocument(self, request, context):
        """同步文档分块请求，使用进程池并行处理多个文件。"""
        logger.info(f"Received sync request for files: {request.file_paths}")
        tasks = [(file_path, request.chunker_types[i], request.chunk_sizes[i], None) for i, file_path in enumerate(request.file_paths)]
        results = self.run_in_process_pool(self.parse_and_chunk_file, tasks, self.max_workers)
        response = document_pb2.ChunkedResponse(
            response_data=[]
        )
        for file_path, chunk, _ in results:
            response.response_data.append(document_pb2.ResponseData(chunks = chunk, file_path = file_path))
        logger.debug(f"Returning sync response: {response}")
        logger.info(f"chunk_over: {request.file_paths}")
        return response

    def AsyncChunkDocument(self, request, context):
        """异步文档分块请求，使用进程池并行处理多个文件。"""
        logger.info(f"Received async request for files: {request.file_paths}")
        logger.info(f"Request file paths: {request.file_paths}")
        logger.info(f"Request chunker types: {list(request.chunker_types)}")
        logger.info(f"Request chunk sizes: {list(request.chunk_sizes)}")
        logger.info(f"Request callback data: {request.callback_data}")
        tasks = [
            (file_path, request.chunker_types[i], request.chunk_sizes[i],request.callback_data[0])
            for i, file_path in enumerate(request.file_paths)
        ]
        logger.info(f"Prepared tasks: {tasks}")

        def background_process(tasks, context):
            with multiprocessing.Pool(processes=self.max_workers) as pool:
                async_results = [pool.apply_async(self.parse_and_chunk_file, task) for task in tasks]
                all_chunks = []
                for async_result in async_results:
                    try:
                        result = async_result.get(timeout=60)
                        (file_path, chunks, callback_data) = result
                        all_chunks.extend(chunks)
                    except Exception as e:
                        logger.error(f"Error processing task: {e}")
                        all_chunks.append(f"Error: {str(e)}")
                
                response = document_pb2.ChunkedResponse(
                    file_paths=request.file_paths,
                    chunks=[str(chunk) for chunk in all_chunks],
                    callback_data=str(request.callback_data[0])
                )
                return response

        def run_background_process():
            response = background_process(tasks, context)
            context.set_code(grpc.StatusCode.OK)
            context.set_details(json.dumps({"status": "completed"}))
            return response
        response_future = self.async_cache.cache_result(run_background_process, key=request.file_paths)
        return response_future.result()

    def run_in_process_pool(self, target, tasks, max_workers):
        with multiprocessing.Pool(processes=max_workers) as pool:
            async_results = [pool.apply_async(target, task) for task in tasks]
            return [async_result.get() for async_result in async_results]

def serve():
    uid = os.getuid()
    home_directory = os.path.expanduser("~")
    print(f"Home directory: {home_directory}")
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10),
                options=[('grpc.max_send_message_length', 1024 * 1024 * 1024),
                        ('grpc.max_receive_message_length', 1024 * 1024 * 1024)])
    document_pb2_grpc.add_DocumentServiceServicer_to_server(DocumentService(), server)
    unix_socket_path_sync = '/tmp/document_service_sync_' + str(uid) +'.sock'
    unix_socket_path_async = '/tmp/document_service_async_' + str(uid) +'.sock'
    try:
        os.unlink(unix_socket_path_sync)
    except OSError:
        if os.path.exists(unix_socket_path_sync):
            raise
    
    try:
        os.unlink(unix_socket_path_async)
    except OSError:
        if os.path.exists(unix_socket_path_async):
            raise
    server.add_insecure_port(f'unix://{unix_socket_path_sync}')
    server.add_insecure_port(f'unix://{unix_socket_path_async}')
    server.start()
    logger.info("Server started and listening on Unix domain sockets.")
    server.wait_for_termination()

if __name__ == '__main__':
    serve()
