import os
import copy
import logging
import pandas as pd
from tqdm.asyncio import tqdm_asyncio
from pathlib import Path
from datetime import datetime
from LLM_call import LLMModel
from LLM_TM import LLMTM
from LLM_CG import LLMCG, CodeRunner
from Evaluator import Evaluator
from PR_predictor import PassRatePredictor
from collections import defaultdict
from typing import Dict, List
from tqdm import tqdm

# 创建Logger实例
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

# 定义日志格式
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')

# 输出到文件的Handler
log_dir = Path("log")
log_dir.mkdir(exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
log_filename = f"{timestamp}.log"
log_path = log_dir / log_filename
file_handler = logging.FileHandler(log_path)
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(formatter)

# 输出到控制台的Handler
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)
console_handler.setFormatter(formatter)

# 将Handler添加到Logger
logging.basicConfig(
    level=logging.DEBUG,
    datefmt="%Y-%m-%d %H:%M:%S",
    handlers=[
        file_handler,   # 文件输出
        console_handler # 控制台输出
    ]
)


class LCDP():
    def __init__(self, api_key, model="gpt-3.5-turbo", max_workers=5, ignore_advice=False, use_pr_predictor=True, use_web_search=False):
        self.llm_model = LLMModel(api_key, model, use_web_search=use_web_search)
        self.code_runner = CodeRunner(max_workers=max_workers)
        self.use_pr_predictor = use_pr_predictor
        if use_pr_predictor:
            self.pass_rate_predictor = PassRatePredictor()
        else:
            self.pass_rate_predictor = None
        self.evaluator = Evaluator(self.pass_rate_predictor)
        self.task_description = None
        self.current_plan = None
        self.test_weights = {}
        self.test_cases = {}
        self.test_timeout = None
        self.ignore_advice = ignore_advice

    def initialize(self):
        if self.use_pr_predictor:
            self.pass_rate_predictor = PassRatePredictor()
            self.evaluator = Evaluator(self.pass_rate_predictor)
        else:
            self.pass_rate_predictor = None
        self.task_description = None
        self.current_plan = None
        self.test_weights = {}
        self.test_cases = {}
        self.test_timeout = None

    async def run(self, 
                  task_description, 
                  max_iterations=3, 
                  stop_t=0.8, 
                  example_dataset=None,
                  num_plans=3, 
                  num_tests=5, 
                  num_codes=5, 
                  refine_rounds=3, 
                  use_pass_rate_for_train=False, 
                  test_timeout=None, 
                  min_tests=15, 
                  max_tests=20, 
                  use_async_generation=True, 
                  knowledge_refine=False,
                  best_only=False, 
                  error_test_num=3, 
                  use_data_format_extract=[], 
                  record_all_results=False, 
                  forced_test_cases=None, 
                  prompt_only=False):
        
        self.initialize()

        # Initialize LLM Task Manager
        self.llmtm = LLMTM(self.llm_model)
        self.llmcg = LLMCG(self.llm_model)

        if knowledge_refine:
            # refine the task description by adding extra explanation for expert knowledge in task description
            # TODO
            print()
            self.task_description = self.llmtm.task_knowledge_refinement(task_description)

        




        