import evoagentx.workflow.operators as operator
import examples.aflow.scicode.optimized.round_6.prompt as prompt_custom
from evoagentx.models.model_configs import LLMConfig
from evoagentx.benchmark.benchmark import Benchmark
from evoagentx.models.model_utils import create_llm_instance

class Workflow:
    
    def __init__(
        self,
        name: str,
        llm_config: LLMConfig,
        benchmark: Benchmark
    ):
        self.name = name
        self.llm = create_llm_instance(llm_config)
        self.benchmark = benchmark 
        self.custom = operator.Custom(self.llm)
        self.custom_code_generate = operator.CustomCodeGenerate(self.llm)
        self.test = operator.Test(self.llm)  # Added Test operator for solution validation
        self.sc_ensemble = operator.ScEnsemble(self.llm)  # Added ScEnsemble operator for solution selection

    async def __call__(self, problem: str, entry_point: str):
        """
        Implementation of the workflow
        Custom operator to generate anything you want.
        Generate code and validate it using the Test operator, then refine using the ScEnsemble operator for the final selection.
        """
        solution = await self.custom_code_generate(problem=problem, entry_point=entry_point, instruction=prompt_custom.GENERATE_PYTHON_CODE_PROMPT) 
        validation = await self.test(problem=problem, solution=solution['response'], entry_point=entry_point, benchmark=self.benchmark)
        
        if not validation['result']:  # If the initial solution fails
            refined_solution = await self.custom(input=problem, instruction=prompt_custom.GENERATE_PYTHON_CODE_PROMPT)  # Regenerate solution
            validation = await self.test(problem=problem, solution=refined_solution['response'], entry_point=entry_point, benchmark=self.benchmark)
        
        # If the validation passes, we can proceed to ensemble solutions
        if validation['result']:
            return self.sc_ensemble(solutions=[solution['response'], refined_solution['response']], problem=problem)['response']  # Select the best solution
        return validation['solution']  # Return the current solution if still failing