from typing import List, Union
from pprint import pprint
from multiprocessing import Queue, Process
from pathlib import Path
from logging import getLogger
from itertools import count
import time
import os

from tqdm import tqdm
from langchain_openai import ChatOpenAI
from langchain.schema.messages import HumanMessage, SystemMessage

from .utils import pickle_bobj, get_saving_filename_safely


logger = getLogger(__name__)


def process_chunk_element(i, queue, item, model_name, temperature):
    chat = ChatOpenAI(model=model_name, temperature=temperature)
    res = chat.invoke(item)
    queue.put((i, res))


def process_chunk(chunk, model_name, temperature, timeout=13):
    processes = []
    output_queue = Queue()
    results = [None] * len(chunk)  # Pre-allocate list for results

    for i, item in enumerate(chunk):
        p = Process(target=process_chunk_element, args=(i, output_queue, item, model_name, temperature))
        processes.append(p)
        p.start()
        time.sleep(0.2)  # restrict too dense api calling

    start_time = time.time()
    completed = 0

    while completed < len(processes) and time.time() - start_time < timeout:
        if not output_queue.empty():
            index, result = output_queue.get()
            results[index] = result
            completed += 1

    # Terminate any remaining processes
    for p in processes:
        if p.is_alive():
            p.terminate()
            p.join()

    return results


def batched_multiprocess_auto_retry(
        items,
        model_name,
        temperature,
        chunk_size,
        timeout_each,
        sleep_between_chunk,
        pkl_path,
        verbose=False):
    """returns list of chatgpt output string

    timeout-ed output be None"""
    pkl_path = get_saving_filename_safely(pkl_path) if pkl_path else None  # if pkl_path, result saved.

    outputs = [None] * len(items)
    c = count()
    while not all(outputs):
        retry_count = next(c)
        if retry_count > 0:
            logger.info("There exists failed API call. Retrying...")
        # display a warning message if the global retry count exceeds 5.
        if retry_count >= 5:
            logger.warning("Retry count has exceeds 5. The process may be stuck due to an unresponsive item.")

        # printing remained queries if the number of remained queries is small
        num_of_remains = outputs.count(None)
        print(f"num of remains: {num_of_remains}") if verbose else ...
        if verbose and num_of_remains <= chunk_size:
            pprint(f"printing remains...:\n{[items[i][1].content for i, o in enumerate(outputs) if o is None]}")

        remain_inputs = [(i, item) for i, item in enumerate(items) if outputs[i] is None]  # store failed item indices
        remain_indices, remain_items = list(zip(*remain_inputs))
        chunks = [remain_items[i:i + chunk_size] for i in range(0, len(remain_items), chunk_size)]  # re-chunk remains

        for i, chunk in enumerate(tqdm(chunks, "Batches")):  # tqdm num is the num of chunks
            results = process_chunk(chunk, model_name, temperature, timeout_each)
            results = list(map(lambda x: x.content if x else None, results))
            for j, result in enumerate(results):
                outputs[remain_indices[i * chunk_size + j]] = result

            # save the outputs which may be incomplete
            pickle_bobj({'prompts': items, 'completions': outputs}, pkl_path) if pkl_path else None

            time.sleep(sleep_between_chunk) if not all(outputs) else ...
    return outputs


def call_chatgpt(
        human_message: List[str],
        system_message: Union[str, List[str]] = "You're a helpful assistant",
        model_name: str = "gpt-3.5-turbo",
        temperature: float = 0.0,
        chunk_size: int = 20,
        timeout_each: int = 60,
        sleep_between_chunk: int = 10,
        pkl_path: Union[Path, str] = None,
        verbose: bool = False) -> List[str]:
    """call batched chatgpt api, and automatically save the responses.

    if pkl_path is not None, then this function automatically save the results to pkl_path.
    :param human_message: list of human message
    :param system_message: list of system prompt. It can be a str or a list of str that has same length to human_message
    :param model_name: ChatGPT API name ("gpt-4-1106-preview")
    :param temperature: Controls randomness of the generated text
    :param chunk_size: The number of examples which send in one batch
    :param timeout_each: API call timeout
    :param sleep_between_chunk: sleep time between batches
    :param pkl_path: Chatgpt output will be saved to this path
    :param verbose: If true, the debugging message will be printed.

    :return: list of the answers from chatgpt
    """
    assert isinstance(system_message, list) or isinstance(system_message, str)
    assert isinstance(human_message, list)
    if isinstance(system_message, str):
        system_message = [system_message] * len(human_message)
    assert len(system_message) == len(human_message)
    assert os.environ['OPENAI_API_KEY']

    messages_list = [[
        SystemMessage(content=system_message[i]),
        HumanMessage(
            content=human_message[i]
        )
    ] for i in range(len(system_message))]

    resp = batched_multiprocess_auto_retry(
        messages_list, model_name, temperature, chunk_size, timeout_each, sleep_between_chunk, pkl_path, verbose)
    return resp