import logging
from openai import OpenAI
from tqdm import tqdm
from collections import defaultdict
import traceback
import httpx

from backend.utils.data_process import split_to_file_diff, split_to_section
from backend.section_infer_helper.base_helper import BaseHelper


logger = logging.getLogger(__name__)


class OnlineLLMHelper(BaseHelper):

    MAX_LENGTH = 4096
    MAX_NEW_TOKENS = 16
    
    PREDEF_MODEL = ["gpt-3.5-turbo", "deepseek-chat", "qwen-coder-plus", "gpt-4-turbo", "gpt-4o", "gemini-1.5-pro-latest", "claude-3-5-sonnet-20241022"]
    
    MODEL_CONFIGS = defaultdict(lambda: {
        "supported_languages": ["C", "C++", "Java", "Python"],
    })
    
    SYSTEM_PROMPT = "You are an expert in code vulnerability and patch fixes."
    
    def generate_instruction(language, file_name, patch, section, message = None):
        instruction =  "[TASK]\nHere is a patch in {} language and a section of this patch for a source code file with path {}. Determine if the patch section fixes any software vulnerabilities. Output 'yes' or 'no' and do not output any other text.\n".format(language, file_name)
        instruction += "[Patch]\n{}\n".format(patch)
        instruction += "[A section of this patch]\n{}\n".format(section)
        if message is not None and message != "":
            instruction += "[Message of the Patch]\n{}\n".format(message)
        
        return instruction


    def __init__(self):
        self.model_name = None
        self.url = None
        self.key = None


    def generate_message(filename, patch, section, patch_message = None):
        ext = filename.split(".")[-1]
        language = BaseHelper._get_lang_by_ext(ext)
        user_message = OnlineLLMHelper.generate_instruction(language, filename, patch, section, patch_message)
        user_message = user_message.split(" ")
        user_message = user_message[:OnlineLLMHelper.MAX_LENGTH]
        user_message = " ".join(user_message)
        messages = [
            {
                "role": "system",
                "content": OnlineLLMHelper.SYSTEM_PROMPT
            },
            {
                "role": "user",
                "content": user_message
            }
        ]
        return messages
    

    def load_model(self, model_name, url, api_key):
        self.model_name = model_name

        self.openai_client = OpenAI(
            base_url = url,
            api_key = api_key,
            timeout=httpx.Timeout(15.0)
        )
    
    
    def infer(self, diff_code, message = None, batch_size=1):
        if self.model_name is None:
            raise RuntimeError("Model is not loaded")
        
        results = {}
        input_list = []
        file_diff_list = split_to_file_diff(diff_code, BaseHelper._get_lang_ext(OnlineLLMHelper.MODEL_CONFIGS[self.model_name]["supported_languages"]))
        for file_a, _, file_diff in file_diff_list:
            sections = split_to_section(file_diff)
            file_name = file_a.removeprefix("a/")
            results[file_name] = []
            for section in sections:
                input_list.append(BaseHelper.InputData(file_name, section, section, message))

        input_prompt, output_text, output_prob = self.do_infer(input_list, batch_size)
        assert len(input_list) == len(input_prompt) == len(output_text) == len(output_prob)
        for i in range(len(input_list)):
            file_name = input_list[i].filename
            section = input_list[i].section
            output_text_i = output_text[i].lower()
            output_prob_i = output_prob[i]
            results[file_name].append({
                "section": section,
                "predict": -1 if output_text_i == "error" else 1 if "yes" in output_text_i else 0,
                "conf": output_prob_i
            })

        return results
        
    
    def do_infer(self, input_list, batch_size = 1):
        input_prompt = []
        for input_data in input_list:
            input_prompt.append(OnlineLLMHelper.generate_message(input_data.filename, input_data.patch, input_data.section, input_data.patch_msg))
        
        if len(input_prompt) > 0:
            logger.info("Example input prompt: %s", input_prompt[0])
        output_text = []
        for prompt, input_data in tqdm(zip(input_prompt, input_list), desc="Inferencing", unit = "section", total = len(input_prompt)):
            try:
                response = self.openai_client.chat.completions.create(
                    messages = prompt,
                    model = self.model_name,
                    max_completion_tokens = OnlineLLMHelper.MAX_NEW_TOKENS
                )
                output_text.append(response.choices[0].message.content)
            except KeyboardInterrupt:
                logging.error("KeyboardInterrupted")
                break
            except Exception as e:
                logger.error(f"Error: {e}")
                logger.error(f"Error inferencing: {input_data.filename} - {input_data.section}")
                logger.error(traceback.format_exc())
                
                output_text.append("error")
                continue
                
                # break
        
        output_prob = [1.0] * len(output_text)
        return input_prompt, output_text, output_prob


online_llm_helper = OnlineLLMHelper()