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

from crewai.flow.flow import start, listen, and_
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.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.globals.identification.identification_v1_0 import LeadInfoOutput
from eveai_chat_workers.outputs.globals.spin.spin_v1_0 import SPINOutput
from eveai_chat_workers.outputs.globals.rag.rag_v1_0 import RAGOutput
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: SPIN_SPECIALIST
    type_version: 1.0
    SPIN Specialist Executor class
    """

    def __init__(self, tenant_id, specialist_id, session_id, task_id, **kwargs):
        self.rag_crew = None
        self.spin_crew = None
        self.identification_crew = None
        self.rag_consolidation_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)
        if self.specialist.configuration['tenant_language'] is None:
            self.specialist.configuration['tenant_language'] = self.tenant.language

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

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

    def _config_task_agents(self):
        self._add_task_agent("rag_task", "rag_agent")
        self._add_task_agent("spin_detect_task", "spin_detection_agent")
        self._add_task_agent("spin_questions_task", "spin_sales_specialist_agent")
        self._add_task_agent("identification_detection_task", "identification_agent")
        self._add_task_agent("identification_questions_task", "identification_agent")
        self._add_task_agent("email_lead_drafting_task", "email_content_agent")
        self._add_task_agent("email_lead_engagement_task", "email_engagement_agent")
        self._add_task_agent("email_lead_retrieval_task", "email_engagement_agent")
        self._add_task_agent("rag_consolidation_task", "rag_communication_agent")

    def _config_pydantic_outputs(self):
        self._add_pydantic_output("rag_task", RAGOutput, "rag_output")
        self._add_pydantic_output("spin_questions_task", SPINOutput, "spin_questions")
        self._add_pydantic_output("identification_questions_task", LeadInfoOutput, "lead_identification_questions")
        self._add_pydantic_output("rag_consolidation_task", RAGOutput, "final_output")

    def _instantiate_specialist(self):
        verbose = self.tuning

        rag_agents = [self.rag_agent]
        rag_tasks = [self.rag_task]
        self.rag_crew = EveAICrewAICrew(
            self,
            "Rag Crew",
            agents=rag_agents,
            tasks=rag_tasks,
            verbose=verbose,
        )

        spin_agents = [self.spin_detection_agent, self.spin_sales_specialist_agent]
        spin_tasks = [self.spin_detect_task, self.spin_questions_task]
        self.spin_crew = EveAICrewAICrew(
            self,
            "SPIN Crew",
            agents=spin_agents,
            tasks=spin_tasks,
            verbose=verbose,
        )

        identification_agents = [self.identification_agent]
        identification_tasks = [self.identification_detection_task, self.identification_questions_task]
        self.identification_crew = EveAICrewAICrew(
            self,
            "Identification Crew",
            agents=identification_agents,
            tasks=identification_tasks,
            verbose=verbose,
        )

        consolidation_agents = [self.rag_communication_agent]
        consolidation_tasks = [self.rag_consolidation_task]
        self.rag_consolidation_crew = EveAICrewAICrew(
            self,
            "Rag Consolidation Crew",
            agents=consolidation_agents,
            tasks=consolidation_tasks,
            verbose=verbose,
        )

        self.flow = SPINFlow(
            self,
            self.rag_crew,
            self.spin_crew,
            self.identification_crew,
            self.rag_consolidation_crew
        )

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

        flow_inputs = {
            "language": arguments.language,
            "query": arguments.query,
            "context": formatted_context,
            "history": self.formatted_history,
            "historic_spin": json.dumps(self.latest_spin, indent=2),
            "historic_lead_info": json.dumps(self.latest_lead_info, indent=2),
            "name": self.specialist.configuration.get('name', ''),
            "company": self.specialist.configuration.get('company', ''),
            "products": self.specialist.configuration.get('products', ''),
            "product_information": self.specialist.configuration.get('product_information', ''),
            "engagement_options": self.specialist.configuration.get('engagement_options', ''),
            "tenant_language": self.specialist.configuration.get('tenant_language', ''),
            "nr_of_questions": self.specialist.configuration.get('nr_of_questions', ''),
            "identification": arguments.identification,
        }

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

        flow_state = self.flow.state

        results = SPINSpecialistResult.create_for_type(self.type, self.type_version)
        update_data = {}
        if flow_state.final_output:
            update_data["rag_output"] = flow_state.final_output
        elif flow_state.rag_output:     # Fallback
            update_data["rag_output"] = flow_state.rag_output
        if flow_state.spin:
            update_data["spin"] = flow_state.spin
        if flow_state.lead_info:
            update_data["lead_info"] = flow_state.lead_info

        results = results.model_copy(update=update_data)

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

        return results


class SPINSpecialistInput(BaseModel):
    language: Optional[str] = Field(None, alias="language")
    query: Optional[str] = Field(None, alias="query")
    context: Optional[str] = Field(None, alias="context")
    citations: Optional[List[int]] = Field(None, alias="citations")
    history: Optional[str] = Field(None, alias="history")
    historic_spin: Optional[str] = Field(None, alias="historic_spin")
    historic_lead_info: Optional[str] = Field(None, alias="historic_lead_info")
    name: Optional[str] = Field(None, alias="name")
    company: Optional[str] = Field(None, alias="company")
    products: Optional[str] = Field(None, alias="products")
    product_information: Optional[str] = Field(None, alias="product_information")
    engagement_options: Optional[str] = Field(None, alias="engagement_options")
    tenant_language: Optional[str] = Field(None, alias="tenant_language")
    nr_of_questions: Optional[int] = Field(None, alias="nr_of_questions")
    identification: Optional[str] = Field(None, alias="identification")


class SPINSpecialistResult(SpecialistResult):
    rag_output: Optional[RAGOutput] = Field(None, alias="Rag Output")
    spin: Optional[SPINOutput] = Field(None, alias="Spin Output")
    lead_info: Optional[LeadInfoOutput] = Field(None, alias="Lead Info Output")


class SPINFlowState(EveAIFlowState):
    """Flow state for SPIN specialist that automatically updates from task outputs"""
    input: Optional[SPINSpecialistInput] = None
    rag_output: Optional[RAGOutput] = None
    lead_info: Optional[LeadInfoOutput] = None
    spin: Optional[SPINOutput] = None
    final_output: Optional[RAGOutput] = None


class SPINFlow(EveAICrewAIFlow[SPINFlowState]):
    def __init__(self,
                 specialist_executor: CrewAIBaseSpecialistExecutor,
                 rag_crew: EveAICrewAICrew,
                 spin_crew: EveAICrewAICrew,
                 identification_crew: EveAICrewAICrew,
                 rag_consolidation_crew: EveAICrewAICrew,
                 **kwargs):
        super().__init__(specialist_executor, "SPIN Specialist Flow", **kwargs)
        self.specialist_executor = specialist_executor
        self.rag_crew = rag_crew
        self.spin_crew = spin_crew
        self.identification_crew = identification_crew
        self.rag_consolidation_crew = rag_consolidation_crew
        self.exception_raised = False

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

    @listen(process_inputs)
    async def execute_rag(self):
        inputs = self.state.input.model_dump()
        try:
            crew_output = await self.rag_crew.kickoff_async(inputs=inputs)
            self.specialist_executor.log_tuning("RAG Crew Output", crew_output.model_dump())
            output_pydantic = crew_output.pydantic
            if not output_pydantic:
                raw_json = json.loads(crew_output.raw)
                output_pydantic = RAGOutput.model_validate(raw_json)
            self.state.rag_output = output_pydantic
            return crew_output
        except Exception as e:
            current_app.logger.error(f"CREW rag_crew Kickoff Error: {str(e)}")
            self.exception_raised = True
            raise e

    @listen(process_inputs)
    async def execute_spin(self):
        inputs = self.state.input.model_dump()
        try:
            crew_output = await self.spin_crew.kickoff_async(inputs=inputs)
            current_app.logger.info(f"SPIN Crew Executed, output: {crew_output.model_dump()}")
            self.specialist_executor.log_tuning("Spin Crew Output", crew_output.model_dump())
            output_pydantic = crew_output.pydantic
            if not output_pydantic:
                raw_json = json.loads(crew_output.raw)
                output_pydantic = SPINOutput.model_validate(raw_json)
            self.state.spin = output_pydantic
            return crew_output
        except Exception as e:
            current_app.logger.error(f"CREW spin_crew Kickoff Error: {str(e)}")
            self.exception_raised = True
            raise e

    @listen(process_inputs)
    async def execute_identification(self):
        inputs = self.state.input.model_dump()
        try:
            crew_output = await self.identification_crew.kickoff_async(inputs=inputs)
            self.specialist_executor.log_tuning("Identification Crew Output", crew_output.model_dump())
            output_pydantic = crew_output.pydantic
            if not output_pydantic:
                raw_json = json.loads(crew_output.raw)
                output_pydantic = LeadInfoOutput.model_validate(raw_json)
            self.state.lead_info = output_pydantic
            return crew_output
        except Exception as e:
            current_app.logger.error(f"CREW identification_crew Kickoff Error: {str(e)}")
            self.exception_raised = True
            raise e

    @listen(and_(execute_rag, execute_spin, execute_identification))
    async def consolidate(self):
        inputs = self.state.input.model_dump()
        if self.state.rag_output:
            inputs["prepared_answers"] = self.state.rag_output.answer
        additional_questions = ""
        if self.state.lead_info:
            additional_questions = self.state.lead_info.questions + "\n"
        if self.state.spin:
            additional_questions = additional_questions + self.state.spin.questions
        inputs["additional_questions"] = additional_questions
        try:
            crew_output = await self.rag_consolidation_crew.kickoff_async(inputs=inputs)
            self.specialist_executor.log_tuning("RAG Consolidation Crew Output", crew_output.model_dump())
            output_pydantic = crew_output.pydantic
            if not output_pydantic:
                raw_json = json.loads(crew_output.raw)
                output_pydantic = RAGOutput.model_validate(raw_json)
            self.state.final_output = output_pydantic
            return crew_output
        except Exception as e:
            current_app.logger.error(f"CREW rag_consolidation_crew Kickoff Error: {str(e)}")
            self.exception_raised = True
            raise e

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