from datetime import datetime as dt, timezone as tz
from typing import Optional, List, Dict

import json
import yaml
from crewai.flow.flow import start, listen
from flask import current_app
from pydantic import BaseModel, Field

from common.extensions import db, minio_client
from common.models.interaction import Specialist, EveAIAsset
from common.utils.minio_utils import MIB_CONVERTOR
from common.utils.eveai_exceptions import EveAISpecialistExecutionError
from common.utils.model_logging_utils import set_logging_information
from eveai_chat_workers.definitions.language_level.language_level_v1_0 import LANGUAGE_LEVEL
from eveai_chat_workers.definitions.tone_of_voice.tone_of_voice_v1_0 import TONE_OF_VOICE
from eveai_chat_workers.outputs.traicie.knockout_questions.knockout_questions_v1_0 import KOQuestions, KOQuestion
from eveai_chat_workers.specialists.crewai_base_classes import EveAICrewAICrew, EveAICrewAIFlow, EveAIFlowState
from eveai_chat_workers.specialists.crewai_base_specialist import CrewAIBaseSpecialistExecutor
from eveai_chat_workers.specialists.specialist_typing import SpecialistResult, SpecialistArguments


class SpecialistExecutor(CrewAIBaseSpecialistExecutor):
    """
    type: TRAICIE_KO_INTERVIEW_DEFINITION_SPECIALIST
    type_version: 1.0
    Traicie Selection Specialist Executor class
    """

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

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


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

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

    def _config_task_agents(self):
        self._add_task_agent("traicie_ko_criteria_interview_definition_task", "traicie_hr_bp_agent")

    def _config_pydantic_outputs(self):
        self._add_pydantic_output("traicie_ko_criteria_interview_definition_task", KOQuestions, "ko_questions")

    def _config_state_result_relations(self):
        self._add_state_result_relation("ko_questions")

    def _instantiate_specialist(self):
        verbose = self.tuning

        ko_def_agents = [self.traicie_hr_bp_agent]
        ko_def_tasks = [self.traicie_ko_criteria_interview_definition_task]
        self.ko_def_crew = EveAICrewAICrew(
            self,
            "KO Criteria Interview Definition Crew",
            agents=ko_def_agents,
            tasks=ko_def_tasks,
            verbose=verbose,
        )

        self.flow = KOFlow(
            self,
            self.ko_def_crew
        )

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

        if not self._cached_session.interactions:
            specialist_phase = "initial"
        else:
            specialist_phase = self._cached_session.interactions[-1].specialist_results.get('phase', 'initial')

        results = None

        match specialist_phase:
            case "initial":
                results = self.execute_initial_state(arguments, formatted_context, citations)

        self.log_tuning(f"Traicie KO Criteria Interview Definition Specialist execution ended",
                        {"Results": results.model_dump() if results else "No info"})

        return results


    def execute_initial_state(self, arguments: SpecialistArguments, formatted_context, citations) -> SpecialistResult:
        self.log_tuning("Traicie KO Criteria Interview Definition Specialist initial_state_execution started", {})

        selection_specialist = Specialist.query.get(arguments.specialist_id)
        if not selection_specialist:
            raise EveAISpecialistExecutionError(self.tenant_id, self.specialist_id, self.session_id,
                                                "No selection specialist found")
        if selection_specialist.type != "TRAICIE_SELECTION_SPECIALIST":
            raise EveAISpecialistExecutionError(self.tenant_id, self.specialist_id, self.session_id,
                                                "Specialist is no Selection Specialist")

        ko_competencies = []
        for competency in selection_specialist.configuration.get("competencies", []):
            if competency["is_knockout"] is True and competency["assess"] is True:
                ko_competencies.append({"title": competency["title"], "description": competency["description"]})

        tone_of_voice = selection_specialist.configuration.get('tone_of_voice', 'Professional & Neutral')
        selected_tone_of_voice = next(
            (item for item in TONE_OF_VOICE if item["name"] == tone_of_voice),
            None  # fallback indien niet gevonden
        )
        tone_of_voice_context = f"{selected_tone_of_voice["description"]}"

        language_level = selection_specialist.configuration.get('language_level', 'Standard')
        selected_language_level = next(
            (item for item in LANGUAGE_LEVEL if item["name"] == language_level),
            None
        )
        language_level_context = (f"{selected_language_level['description']}, "
                                  f"corresponding to CEFR level {selected_language_level['cefr_level']}")

        flow_inputs = {
            'name': "Evie",
            'tone_of_voice': tone_of_voice,
            'tone_of_voice_context': tone_of_voice_context,
            'language_level': language_level,
            'language_level_context': language_level_context,
            'ko_criteria': ko_competencies,
        }

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

        new_type = "TRAICIE_KO_CRITERIA_QUESTIONS"

        # Controleer of we een KOQuestions object hebben of een lijst van KOQuestion objecten
        if hasattr(self.flow.state.ko_questions, 'to_json'):
            # Het is een KOQuestions object
            json_str = self.flow.state.ko_questions.to_json()
        elif isinstance(self.flow.state.ko_questions, list):
            # Het is een lijst van KOQuestion objecten
            # Maak een KOQuestions object en gebruik to_json daarop
            ko_questions_obj = KOQuestions.from_question_list(self.flow.state.ko_questions)
            json_str = ko_questions_obj.to_json()
        else:
            # Fallback voor het geval het een onverwacht type is
            current_app.logger.warning(f"Unexpected type for ko_questions: {type(self.flow.state.ko_questions)}")
            ko_questions_data = [q.model_dump() for q in self.flow.state.ko_questions]
            json_str = json.dumps(ko_questions_data, ensure_ascii=False, indent=2)

        try:
            asset = db.session.query(EveAIAsset).filter(
                EveAIAsset.type == new_type,
                EveAIAsset.type_version == "1.0.0",
                EveAIAsset.configuration.is_not(None),
                EveAIAsset.configuration.has_key('specialist_id'),
                EveAIAsset.configuration['specialist_id'].astext.cast(db.Integer) == selection_specialist.id
            ).first()
        except (ValueError, TypeError) as e:
            current_app.logger.warning(f"Error casting specialist_id in asset configuration: {str(e)}")
            asset = None

        if not asset:
            asset = EveAIAsset(
                name=f"KO Criteria Form for specialist {selection_specialist.id}",
                type=new_type,
                type_version="1.0.0",
                system_metadata={
                    "Creator Specialist Type": self.type,
                    "Creator Specialist Type Version": self.type_version,
                    "Creator Specialist ID": self.specialist_id
                },
                configuration={
                    "specialist_id": selection_specialist.id,
                },
            )
            set_logging_information(asset, dt.now(tz=tz.utc))
            asset.last_used_at = asset.created_at
        else:
            asset.last_used_at = dt.now(tz=tz.utc)

        try:
            # Stap 1: Asset aanmaken maar nog niet committen
            db.session.add(asset)
            db.session.flush()  # Geeft ons het ID zonder te committen

            # Stap 2: Upload naar MinIO (kan falen zonder database impact)
            bucket_name, object_name, file_size = minio_client.upload_asset_file(
                tenant_id=self.tenant_id,
                asset_id=asset.id,
                asset_type=new_type,
                file_type="json",
                file_data=json_str
            )

            # Stap 3: Storage metadata toevoegen
            asset.bucket_name = bucket_name
            asset.object_name = object_name
            asset.file_size = file_size / MIB_CONVERTOR
            asset.file_type = "json"

            # Stap 4: Token usage toevoegen
            asset.prompt_tokens = self.ko_def_crew.usage_metrics.prompt_tokens
            asset.completion_tokens = self.ko_def_crew.usage_metrics.completion_tokens

            # Alles in één keer committen
            db.session.commit()

        except Exception as e:
            current_app.logger.error(f"Error creating asset: {str(e)}")
            db.session.rollback()
            # Probeer MinIO cleanup als upload is gelukt maar database commit faalde
            try:
                if 'bucket_name' in locals() and 'object_name' in locals():
                    minio_client.delete_object(bucket_name, object_name)
            except:
                pass  # Log maar ga door met originele exception
            raise EveAISpecialistExecutionError(self.tenant_id, self.specialist_id, self.session_id,
                                                f"Failed to create asset: {str(e)}")

        results = SpecialistResult.create_for_type(self.type, self.type_version,
                                                   answer=f"asset {asset.id} created for specialist {selection_specialist.id}",
                                                   phase="finished",
                                                   asset_id=asset.id,
                                                   )


        return results


class KODefInput(BaseModel):
    name: Optional[str] = Field(None, alias="name")
    tone_of_voice: Optional[str] = Field(None, alias="tone_of_voice")
    tone_of_voice_context: Optional[str] = Field(None, alias="tone_of_voice_context")
    language_level: Optional[str] = Field(None, alias="language_level")
    language_level_context: Optional[str] = Field(None, alias="language_level_context")
    ko_criteria: Optional[List[Dict[str, str]]] = Field(None, alias="ko_criteria")


class KODefResult(SpecialistResult):
    asset_id: Optional[int] = Field(None, alias="asset_id")


class KOFlowState(EveAIFlowState):
    """Flow state for Traicie Role Definition specialist that automatically updates from task outputs"""
    input: Optional[KODefInput] = None
    ko_questions: Optional[List[KOQuestion]] = Field(None, alias="ko_questions")
    phase: Optional[str] = Field(None, alias="phase")


class KOFlow(EveAICrewAIFlow[KOFlowState]):
    def __init__(self,
                 specialist_executor: CrewAIBaseSpecialistExecutor,
                 ko_def_crew: EveAICrewAICrew,
                 **kwargs):
        super().__init__(specialist_executor, "Traicie KO Interview Definiton Specialist Flow", **kwargs)
        self.specialist_executor = specialist_executor
        self.ko_def_crew = ko_def_crew
        self.exception_raised = False

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

    @listen(process_inputs)
    async def execute_ko_def_definition(self):
        inputs = self.state.input.model_dump()
        try:
            crew_output = await self.ko_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.ko_def_crew.tasks:
                if task.name == "traicie_ko_criteria_interview_definition_task":
                    # update["competencies"] = task.output.pydantic.competencies
                    self.state.ko_questions = task.output.pydantic.ko_questions
            # crew_output.pydantic = crew_output.pydantic.model_copy(update=update)
            self.state.phase = "personal_contact_data"
            return crew_output
        except Exception as e:
            current_app.logger.error(f"CREW execute_ko_def Kickoff Error: {str(e)}")
            self.exception_raised = True
            raise e

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