import pandas as pd
from typing import List, Dict, Callable, Optional
from tqdm import tqdm

from dataflow.utils.registry import OPERATOR_REGISTRY
from dataflow import get_logger
from dataflow.utils.storage import DataFlowStorage
from dataflow.core import OperatorABC

@OPERATOR_REGISTRY.register()
class CodeAutoGeneratedSampleEvaluator(OperatorABC):
    """
    CodeAutoGeneratedSampleEvaluator evaluates code samples based on auto-generation
    markers to provide scores for filtering auto-generated code files.
    """

    # Common markers for auto-generated code
    AUTOGEN_PATTERNS = [
        "auto-generated", "autogenerated", "automatically generated",
        "generated automatically", "this file is generated"
    ]

    def __init__(self, is_generated_func: Optional[Callable[[], bool]] = None):
        """
        Initialize the operator and set up the logger.
        
        Args:
            is_generated_func: Optional external detection function
        """
        self.logger = get_logger()
        self.logger.info(f'Initializing {self.__class__.__name__}...')
        self.score_name = 'CodeAutoGeneratedScore'
        self.is_generated_func = is_generated_func
        self.logger.info(f'{self.__class__.__name__} initialized.')

    @staticmethod
    def get_desc(lang: str = "zh"):
        """
        Provide operator functionality description and parameter documentation.
        """
        if lang == "zh":
            return (
                "基于自动生成标记评估代码样本，检测文件头部的自动生成标记。\n\n"
                "评估指标：\n"
                "- CodeAutoGeneratedMarkerCount: 检测到的自动生成标记数量\n"
                "- CodeAutoGeneratedScore: 综合自动生成得分 (0-1，1表示非自动生成)\n\n"
                "输入要求：需要包含'lines'列\n\n"
                "输出参数：\n"
                "- CodeAutoGeneratedMarkerCount: 检测到的自动生成标记数量\n"
                "- CodeAutoGeneratedScore: 综合自动生成得分"
            )
        else:
            return (
                "Evaluate code samples based on auto-generation markers, detecting auto-generation markers in file headers.\n\n"
                "Evaluation Metrics:\n"
                "- CodeAutoGeneratedMarkerCount: Number of detected auto-generation markers\n"
                "- CodeAutoGeneratedScore: Comprehensive auto-generation score (0-1, 1 means not auto-generated)\n\n"
                "Input Requirement: Requires 'lines' column\n\n"
                "Output Parameters:\n"
                "- CodeAutoGeneratedMarkerCount: Number of detected auto-generation markers\n"
                "- CodeAutoGeneratedScore: Comprehensive auto-generation score"
            )

    def _score_func(self, sample):
        """
        Calculate auto-generation scores for a single code sample.
        
        Args:
            sample: Dictionary containing 'lines' key
            
        Returns:
            Dictionary containing auto-generation scores
        """
        lines = sample.get('lines', [])
        
        # Check external function first
        marker_count = 0
        if self.is_generated_func and self.is_generated_func():
            marker_count = 1
        
        # Check first 5 lines for markers
        first_5 = "\n".join(lines[:5]).lower()
        for pat in self.AUTOGEN_PATTERNS:
            if pat in first_5:
                marker_count += 1
        
        # Calculate comprehensive score (0-1)
        score = 1.0 if marker_count == 0 else 0.0
        
        return {
            'CodeAutoGeneratedMarkerCount': marker_count,
            'CodeAutoGeneratedScore': score
        }

    def eval(self, dataframe, input_key):
        """
        Evaluate auto-generation markers for all samples in the dataframe.
        
        Args:
            dataframe: Input DataFrame
            input_key: Key containing the sample data
            
        Returns:
            List of score dictionaries
        """
        scores_list = []
        self.logger.info(f"Evaluating {self.score_name}...")
        
        for _, row in dataframe.iterrows():
            sample = row[input_key] if isinstance(row[input_key], dict) else {"lines": row[input_key]}
            scores = self._score_func(sample)
            scores_list.append(scores)
        
        self.logger.info("Evaluation complete!")
        return scores_list

    def run(self, storage: DataFlowStorage, input_key: str):
        """
        Execute auto-generation marker evaluation operation.
        
        Args:
            storage: Data storage object
            input_key: Key name for input data
        """
        self.input_key = input_key
        dataframe = storage.read("dataframe")
        self.logger.info("CodeAutoGeneratedScore ready to evaluate.")
        
        scores = self.eval(dataframe, input_key)
        
        # Flatten the nested dictionary of scores into the dataframe
        for idx, score_dict in enumerate(scores):
            for key, value in score_dict.items():
                dataframe.at[idx, key] = value
        
        storage.write(dataframe)
