import asyncio
import json
from os import wait
from typing import Optional, List

from crewai.flow.flow import start, listen, and_
from crewai import Process
from flask import current_app
from gevent import sleep
from pydantic import BaseModel, Field

from common.extensions import cache_manager
from common.models.user import Tenant
from common.utils.business_event_context import current_event
from eveai_chat_workers.outputs.globals.basic_types.list_item import ListItem
from eveai_chat_workers.retrievers.retriever_typing import RetrieverArguments
from eveai_chat_workers.specialists.crewai_base_specialist import CrewAIBaseSpecialistExecutor
from eveai_chat_workers.specialists.specialist_typing import SpecialistResult, SpecialistArguments
from eveai_chat_workers.outputs.traicie.competencies.competencies_v1_0 import BehaviouralCompetencies
from eveai_chat_workers.outputs.traicie.knockout_criteria.knockout_criteria_v1_0 import KnockoutCriteria
from eveai_chat_workers.specialists.crewai_base_classes import EveAICrewAICrew, EveAICrewAIFlow, EveAIFlowState
from common.utils.pydantic_utils import flatten_pydantic_model


class SpecialistExecutor(CrewAIBaseSpecialistExecutor):
    """
    type: TRAICIE_VACANCY_DEFINITION_SPECIALIST
    type_version: 1.0
    Traicie Vacancy Definition Specialist Executor class
    """

    def __init__(self, tenant_id, specialist_id, session_id, task_id, **kwargs):
        self.vac_def_crew = None

        super().__init__(tenant_id, specialist_id, session_id, task_id)

        # Load the Tenant & set language
        self.tenant = Tenant.query.get_or_404(tenant_id)

    @property
    def type(self) -> str:
        return "TRAICIE_VACANCY_DEFINITION_SPECIALIST"

    @property
    def type_version(self) -> str:
        return "1.0"

    def _config_task_agents(self):
        self._add_task_agent("traicie_get_competencies_task", "traicie_hr_bp_agent")
        self._add_task_agent("traicie_get_ko_criteria_task", "traicie_hr_bp_agent")

    def _config_pydantic_outputs(self):
        self._add_pydantic_output("traicie_get_competencies_task", BehaviouralCompetencies, "competencies")
        self._add_pydantic_output("traicie_get_ko_criteria_task", KnockoutCriteria, "criteria")

    def _instantiate_specialist(self):
        verbose = self.tuning

        vac_def_agents = [self.traicie_hr_bp_agent]
        vac_def_tasks = [self.traicie_get_competencies_task, self.traicie_get_ko_criteria_task]
        self.vac_def_crew = EveAICrewAICrew(
            self,
            "Vacancy Definition Crew",
            agents=vac_def_agents,
            tasks=vac_def_tasks,
            verbose=verbose,
        )

        self.flow = VacancyDefinitionFlow(
            self,
            self.vac_def_crew
        )

    def execute(self, arguments: SpecialistArguments, formatted_context, citations) -> SpecialistResult:
        self.log_tuning("Traicie Vacancy Definition Specialist execution started", {})

        flow_inputs = {
            "vacancy_text": arguments.vacancy_text,
        }

        flow_results = self.flow.kickoff(inputs=flow_inputs)

        flow_state = self.flow.state

        results = VacancyDefinitionSpecialistResult.create_for_type(self.type, self.type_version)
        update_data = {}
        if flow_state.competencies:
            results.competencies = flow_state.competencies
        if flow_state.criteria:
            results.criteria = flow_state.criteria

        self.log_tuning(f"Traicie Vacancy Definition Specialist execution ended", {"Results": results.model_dump()})

        return results


class VacancyDefinitionSpecialistInput(BaseModel):
    vacancy_text: Optional[str] = Field(None, alias="vacancy_text")


class VacancyDefinitionSpecialistResult(SpecialistResult):
    competencies: Optional[List[ListItem]] = None
    criteria: Optional[List[ListItem]] = None


class VacancyDefFlowState(EveAIFlowState):
    """Flow state for Traicie Vacancy Definition specialist that automatically updates from task outputs"""
    input: Optional[VacancyDefinitionSpecialistInput] = None
    competencies: Optional[List[ListItem]] = None
    criteria: Optional[List[ListItem]] = None


class VacancyDefinitionFlow(EveAICrewAIFlow[VacancyDefFlowState]):
    def __init__(self,
                 specialist_executor: CrewAIBaseSpecialistExecutor,
                 vac_def_crew: EveAICrewAICrew,
                 **kwargs):
        super().__init__(specialist_executor, "Traicie Vacancy Definition Specialist Flow", **kwargs)
        self.specialist_executor = specialist_executor
        self.vac_def_crew = vac_def_crew
        self.exception_raised = False

    @start()
    def process_inputs(self):
        return ""

    @listen(process_inputs)
    async def execute_vac_def(self):
        inputs = self.state.input.model_dump()
        try:
            current_app.logger.debug("In execute_vac_def")
            crew_output = await self.vac_def_crew.kickoff_async(inputs=inputs)
            # Unfortunately, crew_output will only contain the output of the latest task.
            # As we will only take into account the flow state, we need to ensure both competencies and criteria
            # are copies to the flow state.
            update = {}
            for task in self.vac_def_crew.tasks:
                current_app.logger.debug(f"Task {task.name} output:\n{task.output}")
                if task.name == "traicie_get_competencies_task":
                    # update["competencies"] = task.output.pydantic.competencies
                    self.state.competencies = task.output.pydantic.competencies
                elif task.name == "traicie_get_ko_criteria_task":
                    # update["criteria"] = task.output.pydantic.criteria
                    self.state.criteria = task.output.pydantic.criteria
            # crew_output.pydantic = crew_output.pydantic.model_copy(update=update)
            return crew_output
        except Exception as e:
            current_app.logger.error(f"CREW execute_vac_def Kickoff Error: {str(e)}")
            self.exception_raised = True
            raise e

    async def kickoff_async(self, inputs=None):
        self.state.input = VacancyDefinitionSpecialistInput.model_validate(inputs)
        result = await super().kickoff_async(inputs)
        return self.state
