from crewai import Agent, Crew, Process, Task
from crewai.project import CrewBase, agent, crew, task
from crewai import LLM
from user_journey_service.core.config import settings
from user_journey_service.processors.guardrail_fn import GuardRailFn
from crewai import TaskOutput, Task
from typing import Optional, Tuple, Union
import re
import os
from functools import partial
from crewai_tools import FileReadTool
from dotenv import load_dotenv
from crewai.tools import tool
from user_journey_service.tools.custom_tool import read_defaults




load_dotenv()

@CrewBase
class UserJourney():
    """User Journey Development crew"""
    def __init__(self):
       self.llm = LLM(
          model=f"{settings.PROVIDER}/{settings.LLM1}",  
          temperature=settings.TEMPERATURE,
          api_key=settings.OPENAI_KEY1,
          api_base=settings.ENDPOINT1,
          api_version=settings.API_VERSION1,
          seed=settings.SEED,                # For reproducible results
          timeout=settings.TIMEOUT
	)
       self.llm2 = LLM(
          model=f"{settings.PROVIDER}/{settings.LLM2}",    ###     
          temperature=settings.TEMPERATURE,
          api_key=settings.OPENAI_KEY2,
          api_base=settings.ENDPOINT2,
          api_version=settings.API_VERSION2,
          seed=settings.SEED,                # For reproducible results
          timeout=settings.TIMEOUT
	)

    agents_config = 'config/agents.yaml'
    tasks_config = 'config/tasks.yaml'

    @agent
    def defaults_estimator(self) -> Agent:
        return Agent(
            config=self.agents_config['defaults_estimator'],
            llm=self.llm,
            verbose=True
        )

    @agent
    def researcher(self) -> Agent:
        return Agent(
            config=self.agents_config['researcher'],
            llm=self.llm,
            verbose=True

        )

    @agent
    def restructure(self) -> Agent:
        return Agent(
            config=self.agents_config['restructure'],
            llm=self.llm,
            verbose=True
        )
    
    @agent
    def restructure2(self) -> Agent:
        return Agent(
            config=self.agents_config['restructure'],
            llm=self.llm2,
            verbose=True
        )
    
    @agent
    def content_creature(self) -> Agent:
        return Agent(
            config=self.agents_config['content_creature'],
            llm=self.llm,
            verbose=True
        )
    
    @agent
    def question_generator(self) -> Agent:
        return Agent(
            config=self.agents_config['question_generator'],
            llm=self.llm,
            verbose=True
        )
    
    @agent
    def answer_evaluator(self) -> Agent:
        return Agent(
            config=self.agents_config['answer_evaluator'],
            llm=self.llm,
            verbose=True
        )
    
    @agent
    def final_assessment(self) -> Agent:
        return Agent(
            config=self.agents_config['final_assessment'],
            llm=self.llm,
            verbose=True
        )

    @agent
    def ondemand_learning(self) -> Agent:
        return Agent(
            config=self.agents_config['ondemand_learning'],
            llm=self.llm,
            verbose=True
        )
    @agent
    def structured_learning(self) -> Agent:
        return Agent(
            config=self.agents_config['structured_learning'],
            llm=self.llm,
            verbose=True
        )

    

    def defaults_estimator_task(self,output_file) -> Task:
        return Task(
            config=self.tasks_config['defaults_estimator_task'],
            output_file=output_file,
            guardrail=GuardRailFn.validate_defaults_op()  
        )
    
    def research_task(self,research_file) -> Task:
        return Task(
            config=self.tasks_config['research_task'],
            output_file = research_file
        )

    def restructure_task(self,output_file: str,available_time: str, feedback: str) -> Task:
        if feedback =="First iteration":
            config=self.tasks_config['restructure_task']
        else:
            config=self.tasks_config['feedback_task']
        return Task(
            config=config,
            output_file=output_file,
            # tools=[read_defaults],
            guardrail=GuardRailFn.get_validate_content_guardrail(available_time)               ##self.validate_content
        )
    
    def restructure_task2(self,output_file: str,available_time: str, feedback: str) -> Task:
        if feedback =="First iteration":
            config=self.tasks_config['restructure_task']
        else:
            config=self.tasks_config['feedback_task']
        return Task(
            config=config,
            # context=[self.research_task()],
            output_file=output_file,
            guardrail=GuardRailFn.get_validate_content_guardrail(available_time)               ##self.validate_content
        )
    
    
    
    def content_creature_task(self,output_file) -> Task:
        return Task(
            config=self.tasks_config['content_creature_task'],
            output_file=output_file,
        )
    
    def question_generator_task(self,output_file) -> Task:
        return Task(
            config=self.tasks_config['question_generator_task'],
            output_file=output_file,
        )
    
    def answer_evaluator_task(self) -> Task:
        return Task(
            config=self.tasks_config['answer_evaluator_task'],
        )
    
    def final_assessment_task(self) -> Task:
        return Task(
            config=self.tasks_config['final_assessment_task'],
            
        )

    def ondemand_learning_task(self,output_file) -> Task:
        return Task(
            config=self.tasks_config['ondemand_learning_task'],
            output_file=output_file
        )

    def structured_learning_task(self,output_file) -> Task:
        return Task(
            config=self.tasks_config['structured_learning_task'],
            output_file=output_file
        )

    

    def defaults_crew(self,output_file):
        """Creates the Content Development crew"""
        print("Inside defaults crew....")
        return Crew(
            agents=[
                self.defaults_estimator(),
            ],
            tasks=[
                self.defaults_estimator_task(output_file),
            ],
            process=Process.sequential,
            verbose=True,
        )
    
    def researcher_crew(self,research_file):
        return Crew(
            agents=[
                self.researcher(),
            ],
            tasks=[
                self.research_task(research_file),
            ],
            process=Process.sequential,
            verbose=True,
        )


    def restructure_crew1(self,output_file_1: str,available_time: str, feedback: str) -> Crew:
        """Creates the User Journey Development crew"""
        return Crew(
            agents=[
                self.restructure(),
            ],
            tasks=[
                self.restructure_task(output_file_1,available_time,feedback),
            ],
            process=Process.sequential,
            verbose=True,
        )
    
    def restructure_crew2(self, output_file_2: str, available_time: str, feedback: str) -> Crew:
        """Creates the User Journey Development crew"""
        return Crew(
            agents=[
                self.restructure2()
            ],
            tasks=[
                self.restructure_task2(output_file_2,available_time,feedback),
            ],
            process=Process.sequential,
            verbose=True,
        )
    
    
    
    def second_stage_crew(self,output_file: str):
        """Creates the Content Development crew"""
        return Crew(
            agents=[
                self.content_creature(),
            ],
            tasks=[
                self.content_creature_task(output_file),
            ],
            process=Process.sequential,
            verbose=True,
        )
    
    def qa_gen_crew(self,output_file: str):
        """Creates the Content Development crew"""
        return Crew(
            agents=[
                self.question_generator(),
            ],
            tasks=[
                self.question_generator_task(output_file),
            ],
            process=Process.sequential,
            verbose=True,
        )
    def evaluator_crew(self):
        """Creates the Content Development crew"""
        return Crew(
            agents=[
                self.answer_evaluator(),
            ],
            tasks=[
                self.answer_evaluator_task(),
            ],
            process=Process.sequential,
            verbose=True,
        )
    def assessment_crew(self):
        """Creates the Content Development crew"""
        return Crew(
            agents=[
                self.final_assessment(),
            ],
            tasks=[
                self.final_assessment_task(),
            ],
            process=Process.sequential,
            verbose=True,
        )

    def ondemand_crew(self,output_file):
        """Creates the Content Development crew"""
        return Crew(
            agents=[
                self.ondemand_learning(),
            ],
            tasks=[
                self.ondemand_learning_task(output_file),
            ],
            process=Process.sequential,
            verbose=True,
        )

    def structured_crew(self,output_file):
        """Creates the Content Development crew"""
        return Crew(
            agents=[
                self.structured_learning(),
            ],
            tasks=[
                self.structured_learning_task(output_file),
            ],
            process=Process.sequential,
            verbose=True,
        )