import logging
import argparse
import os
import pandas as pd
from templates import PROMPT_STRATEGY

try:
    from aphrodite import LLM, SamplingParams

    print("- Using aphrodite-engine")

except ImportError:
    from vllm import LLM, SamplingParams

    print("- Using vLLM")

# 로깅 설정
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

parser = argparse.ArgumentParser()
parser.add_argument("-g", "--gpu_devices", help=" : CUDA_VISIBLE_DEVICES", default="0")
parser.add_argument(
    "-m",
    "--model",
    help=" : write huggingface model name to evaluate",
    default="LDCC/Chat-Mistral-Nemo-12B-32k",
)
parser.add_argument(
    "-ml", "--model_len", help=" : Maximum Model Length", default=4096, type=int
)
args = parser.parse_args()

logger.info(f"Args - {args}")

os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_devices
gpu_counts = len(args.gpu_devices.split(","))

# LLM 초기화
logger.info(f"Initializing LLM with model: {args.model}")
llm = LLM(
    model=args.model,
    tensor_parallel_size=gpu_counts,
    max_model_len=args.model_len,
    gpu_memory_utilization=0.8,
    trust_remote_code=True, 
)
logger.info("LLM initialized successfully")

sampling_params = SamplingParams(
    temperature=0,
    skip_special_tokens=True,
    max_tokens=args.model_len,
    stop=[
        "<|endoftext|>",
        "[INST]",
        "[/INST]",
        "<|im_end|>",
        "<|end|>",
        "<|eot_id|>",
        "<end_of_turn>",
        "<eos>",
    ],
)

# chat_temlate가 없다면 default로 세팅하는 과정
tokenizer = llm.llm_engine.tokenizer.tokenizer

if tokenizer.chat_template is None:
    logger.info("chat template가 없으므로 default로 설정")
    default_chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
    tokenizer.chat_template = default_chat_template

# 문제 로드
logger.info("Loading questions from questions.jsonl")
df_questions = pd.read_json(
    "questions.jsonl", orient="records", encoding="utf-8-sig", lines=True
)
logger.info(f"Loaded {len(df_questions)} questions")

if not os.path.exists("./generated/" + args.model):
    os.makedirs("./generated/" + args.model)

for strategy_name, prompts in PROMPT_STRATEGY.items():
    logger.info(f"Processing strategy: {strategy_name}")

    def format_single_turn_question(question):
        return tokenizer.apply_chat_template(
            prompts + [{"role": "user", "content": question[0]}],
            tokenize=False,
            add_generation_prompt=True,
        )

    single_turn_questions = df_questions["questions"].map(format_single_turn_question)
    print(single_turn_questions.iloc[0])
    
    # 단일 턴 질문 처리
    logger.info("Generating single-turn outputs")
    single_turn_outputs = [
        output.outputs[0].text.strip()
        for output in llm.generate(single_turn_questions, sampling_params)
    ]
    logger.info(f"Generated {len(single_turn_outputs)} single-turn outputs")

    def format_double_turn_question(question, single_turn_output):
        return tokenizer.apply_chat_template(
            prompts
            + [
                {"role": "user", "content": question[0]},
                {"role": "assistant", "content": single_turn_output},
                {"role": "user", "content": question[1]},
            ],
            tokenize=False,
            add_generation_prompt=True,
        )

    multi_turn_questions = df_questions[["questions", "id"]].apply(
        lambda x: format_double_turn_question(
            x["questions"], single_turn_outputs[x["id"] - 1]
        ),
        axis=1,
    )
    
    # 멀티 턴 질문 처리
    logger.info("Generating multi-turn outputs")
    multi_turn_outputs = [
        output.outputs[0].text.strip()
        for output in llm.generate(multi_turn_questions, sampling_params)
    ]
    logger.info(f"Generated {len(multi_turn_outputs)} multi-turn outputs")

    df_output = pd.DataFrame(
        {
            "id": df_questions["id"],
            "category": df_questions["category"],
            "questions": df_questions["questions"],
            "outputs": list(zip(single_turn_outputs, multi_turn_outputs)),
            "references": df_questions["references"],
        }
    )
    
    # 결과 저장
    output_file = f"./generated/{args.model}/{strategy_name}.jsonl"
    logger.info(f"Saving results to {output_file}")
    df_output.to_json(
        output_file,
        orient="records",
        lines=True,
        force_ascii=False,
    )
    logger.info(f"Results saved for strategy: {strategy_name}")

logger.info("Generation process completed")