import pandas as pd
import numpy as np
from typing import List, Callable, Optional

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

@OPERATOR_REGISTRY.register()
class CodeAutoGeneratedFilter(OperatorABC):
    """
    CodeAutoGeneratedFilter filters auto-generated code files using
    CodeAutoGeneratedSampleEvaluator scores to ensure only human-written code is retained.
    """

    def __init__(self, min_score: float = 1.0, max_score: float = 1.0, is_generated_func: Optional[Callable[[], bool]] = None):
        """
        Initialize the operator with evaluator and thresholds.
        
        Args:
            min_score: Minimum auto-generation score threshold
            max_score: Maximum auto-generation score threshold
            is_generated_func: Optional external detection function
        """
        self.min_score = min_score
        self.max_score = max_score
        self.logger = get_logger()
        self.logger.info(f"Initializing {self.__class__.__name__} with min_score: {self.min_score} and max_score: {self.max_score}...")
        self.scorer = CodeAutoGeneratedSampleEvaluator(is_generated_func)

    @staticmethod
    def get_desc(lang: str = "en"):
        """
        Provide operator functionality description and parameter documentation.
        """
        if lang == "zh":
            return (
                "基于CodeAutoGeneratedSampleEvaluator的得分过滤自动生成的代码文件，确保只保留人工编写的代码。\n\n"
                "评估指标：\n"
                "- 自动生成标记数量：检测文件前5行中的自动生成标记\n"
                "- 检测标记：'auto-generated', 'autogenerated', 'automatically generated'等\n"
                "- 综合自动生成得分：0-1，1表示非自动生成\n"
                "- 支持外部检测函数进行额外验证\n\n"
                "输入参数：\n"
                "- input_key: 输入字段名（需要包含'lines'列）\n"
                "- output_key: 输出标签字段名 (默认: 'auto_generated_filter_label')\n"
                "- min_score: 最小自动生成得分阈值 (默认: 1.0)\n"
                "- max_score: 最大自动生成得分阈值 (默认: 1.0)\n"
                "- is_generated_func: 可选的外部检测函数\n\n"
                "输出参数：\n"
                "- 过滤后的DataFrame，仅保留自动生成得分在指定范围内的代码样本\n"
                "- 返回包含自动生成得分标签字段名的列表"
            )
        else:
            return (
                "Filter auto-generated code files using scores from CodeAutoGeneratedSampleEvaluator to ensure only human-written code is retained.\n\n"
                "Evaluation Metrics:\n"
                "- Auto-generation marker count: Detect markers in first 5 lines\n"
                "- Detect markers: 'auto-generated', 'autogenerated', 'automatically generated', etc.\n"
                "- Comprehensive auto-generation score: 0-1, 1 means not auto-generated\n"
                "- Support external detection functions for additional validation\n\n"
                "Input Parameters:\n"
                "- input_key: Input field name (requires 'lines' column)\n"
                "- output_key: Output label field name (default: 'auto_generated_filter_label')\n"
                "- min_score: Minimum auto-generation score threshold (default: 1.0)\n"
                "- max_score: Maximum auto-generation score threshold (default: 1.0)\n"
                "- is_generated_func: Optional external detection function\n\n"
                "Output Parameters:\n"
                "- Filtered DataFrame containing only code samples with auto-generation scores within specified range\n"
                "- List containing auto-generation score label field name"
            )

    def _validate_dataframe(self, dataframe: pd.DataFrame):
        """
        Validate DataFrame to ensure required columns exist.
        """
        required_keys = [self.input_lines_key]
        missing = [k for k in required_keys if k not in dataframe.columns]

        if missing:
            raise ValueError(f"AutogenFilter missing required columns: {missing}")

    def run(
        self,
        storage: DataFlowStorage,
        input_key: str,
        output_key: str = "auto_generated_filter_label"
    ) -> List[str]:
        """
        Execute auto-generated code detection and filtering using evaluator scores.

        Args:
            storage: Data storage object
            input_key: Field name containing code lines
            output_key: Key name for output label

        Returns:
            List[str]: List containing output key name
        """
        self.input_key = input_key
        self.output_key = output_key
        self.logger.info(f"Running {self.__class__.__name__} with input_key: {self.input_key} and output_key: {self.output_key}...")
        
        dataframe = storage.read("dataframe")
        scores = self.scorer.eval(dataframe, self.input_key)
        
        # Add scores to dataframe
        for idx, score_dict in enumerate(scores):
            for key, value in score_dict.items():
                dataframe.at[idx, key] = value
        
        # Apply filtering based on CodeAutoGeneratedScore
        results = np.ones(len(dataframe), dtype=int)
        score_filter = (dataframe["CodeAutoGeneratedScore"] >= self.min_score) & (dataframe["CodeAutoGeneratedScore"] <= self.max_score)
        nan_filter = np.isnan(dataframe["CodeAutoGeneratedScore"])
        metric_filter = score_filter | nan_filter
        results = results & metric_filter.astype(int)
        
        self.logger.debug(f"Filtered by auto-generated score, {np.sum(results)} data remained")
        dataframe[f"{self.output_key}"] = metric_filter.astype(int)
        
        filtered_dataframe = dataframe[results == 1]
        storage.write(filtered_dataframe)
        self.logger.info(f"Filtering completed. Total records passing filter: {len(filtered_dataframe)}.")
        
        return [self.output_key]