import json
from typing import Dict, Any, Optional, Type, TypeVar, List, Tuple

from crewai.flow.flow import FlowState
from flask import current_app
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

from common.models.interaction import Specialist
from common.utils.business_event_context import current_event
from common.utils.model_utils import get_crewai_llm, create_language_template, get_template
from eveai_chat_workers.retrievers.retriever_typing import RetrieverArguments
from eveai_chat_workers.specialists.crewai_base_classes import EveAICrewAIAgent, EveAICrewAITask
from crewai.tools import BaseTool
from abc import ABC, abstractmethod

from pydantic import BaseModel

from common.extensions import cache_manager
from eveai_chat_workers.specialists.base_specialist import BaseSpecialistExecutor
from common.utils.cache.crewai_configuration import (
    ProcessedAgentConfig, ProcessedTaskConfig, ProcessedToolConfig,
    SpecialistProcessedConfig
)
from eveai_chat_workers.specialists.specialist_typing import SpecialistArguments, SpecialistResult

T = TypeVar('T')  # For generic type hints


class CrewAIBaseSpecialistExecutor(BaseSpecialistExecutor):
    """Base class for all CrewAI-based specialists"""

    def __init__(self, tenant_id: int, specialist_id: int, session_id: str, task_id):
        super().__init__(tenant_id, specialist_id, session_id, task_id)

        self.tuning = self.specialist.tuning
        # Initialize retrievers
        self.retrievers = self._initialize_retrievers()

        # initialize the Flow
        self.flow = None

        # Runtime instances
        self._agents: Dict[str, EveAICrewAIAgent] = {}
        self._tasks: Dict[str, EveAICrewAITask] = {}
        self._tools: Dict[str, BaseTool] = {}

        # Crew configuration
        self._task_agents: Dict[str, str] = {}
        self._task_pydantic_outputs: Dict[str, Type[BaseModel]] = {}
        self._task_state_names: Dict[str, str] = {}

        # State-Result relations (for adding / restoring information to / from history
        self._state_result_relations: Dict[str, str] = {}

        # Process configurations
        self._config = cache_manager.crewai_processed_config_cache.get_specialist_config(tenant_id, specialist_id)
        self._config_task_agents()
        self._config_pydantic_outputs()
        self._instantiate_crew_assets()
        self._instantiate_specialist()
        self._config_state_result_relations()

        # Retrieve history
        self._cached_session = cache_manager.chat_session_cache.get_cached_session(self.session_id)
        self._restore_state_from_history()
        # Format history for the prompt
        self._formatted_history = self._generate_formatted_history()

        self.arguments = None
        self.formatted_context = None
        self.citations = None

    @property
    def formatted_history(self) -> str:
        if not self._formatted_history:
            self._formatted_history = self._generate_formatted_history()
        return self._formatted_history

    def _generate_formatted_history(self) -> str:
        """Generate the formatted history string from cached session interactions."""
        current_app.logger.debug(f"Generating formatted history for {self.session_id}")
        current_app.logger.debug(f"Cached session interactions: {len(self._cached_session.interactions)}")

        return "\n\n".join([
            "\n\n".join([
                f"HUMAN:\n"
                f"{interaction.specialist_arguments['question']}"
                if interaction.specialist_arguments.get('question') and interaction.specialist_arguments.get('question') != "Initialize" else "",
                f"{interaction.specialist_arguments.get('form_values')}"
                if interaction.specialist_arguments.get('form_values') else "",
                f"AI:\n{interaction.specialist_results['answer']}"
                if interaction.specialist_results.get('answer') else ""
            ]).strip()
            for interaction in self._cached_session.interactions
        ])

    def _add_task_agent(self, task_name: str, agent_name: str):
        self._task_agents[task_name.lower()] = agent_name

    @abstractmethod
    def _config_task_agents(self):
        """Configure the task agents by adding task-agent combinations. Use _add_task_agent()
        """
        raise NotImplementedError

    @property
    def task_agents(self) -> Dict[str, str]:
        return self._task_agents

    def _add_pydantic_output(self, task_name: str, output: Type[BaseModel], state_name: str is None):
        self._task_pydantic_outputs[task_name.lower()] = output
        if state_name is not None:
            self._task_state_names[task_name.lower()] = state_name

    @abstractmethod
    def _config_pydantic_outputs(self):
        """Configure the task pydantic outputs by adding task-output combinations. Use _add_pydantic_output()"""
        raise NotImplementedError

    def _add_state_result_relation(self, state_name: str, result_name: str = None):
        """Add a state-result relation to the specialist. This is used to add information to the history
           If result_name is None, the state name is used as the result name. (default behavior)
        """
        if not result_name:
            result_name = state_name
        self._state_result_relations[state_name] = result_name

    def _config_default_state_result_relations(self):
        for default_attribute_name in ['answer', 'form_request', 'phase', 'citations']:
            self._add_state_result_relation(default_attribute_name)

    @abstractmethod
    def _config_state_result_relations(self):
        """Configure the state-result relations by adding state-result combinations. Use _add_state_result_relation()"""
        raise NotImplementedError

    @property
    def task_pydantic_outputs(self):
        return self._task_pydantic_outputs

    @property
    def task_state_names(self):
        return self._task_state_names

    def _instantiate_crew_assets(self):
        self._instantiate_crew_agents()
        self._instantiate_tasks()
        self._instantiate_tools()

    def _instantiate_crew_agents(self):
        for agent in self.specialist.agents:
            agent_config = cache_manager.agents_config_cache.get_config(agent.type, agent.type_version)
            agent_role = agent_config.get('role', '').replace('{custom_role}', agent.role or '')
            agent_role = self._replace_system_variables(agent_role)
            agent_goal = agent_config.get('goal', '').replace('{custom_goal}', agent.goal or '')
            agent_goal = self._replace_system_variables(agent_goal)
            agent_backstory = agent_config.get('backstory', '').replace('{custom_backstory}', agent.backstory or '')
            agent_backstory = self._replace_system_variables(agent_backstory)
            agent_full_model_name = agent_config.get('full_model_name', 'mistral.mistral-large-latest')
            current_app.logger.debug(f"Full model name for {agent.type}: {agent_full_model_name}")
            agent_temperature = agent_config.get('temperature', 0.3)
            llm = get_crewai_llm(agent_full_model_name, agent_temperature)
            if not llm:
                current_app.logger.error(f"No LLM found for {agent_full_model_name}")
                raise Exception(f"No LLM found for {agent_full_model_name}")
            agent_kwargs = {
                "role": agent_role,
                "goal": agent_goal,
                "backstory": agent_backstory,
                "verbose": agent.tuning,
                "llm": llm,
            }
            agent_name = agent.type.lower()
            new_agent = EveAICrewAIAgent(self, agent_name, **agent_kwargs)
            self.log_tuning(f"CrewAI Agent {agent_name} initialized", agent_config)
            self._agents[agent_name] = new_agent

    def _instantiate_tasks(self):
        for task in self.specialist.tasks:
            task_config = cache_manager.tasks_config_cache.get_config(task.type, task.type_version)
            task_description = (task_config.get('task_description', '')
                                .replace('{custom_description}', task.task_description or ''))
            task_description = self._replace_system_variables(task_description)
            task_expected_output = (task_config.get('expected_output', '')
                                    .replace('{custom_expected_output}', task.expected_output or ''))
            # dynamically build the arguments
            task_kwargs = {
                "description": task_description,
                "expected_output": task_expected_output,
                "verbose": task.tuning
            }
            task_name = task.type.lower()
            if task_name in self._task_pydantic_outputs:
                task_kwargs["output_pydantic"] = self._task_pydantic_outputs[task_name]
            if task_name in self._task_agents:
                task_kwargs["agent"] = self._agents[self._task_agents[task_name]]

            # Instantiate the task with dynamic arguments
            new_task = EveAICrewAITask(self, task_name, **task_kwargs)

            # Logging and storing the task
            self.log_tuning(f"CrewAI Task {task_name} initialized", task_config)
            self._tasks[task_name] = new_task

    def _instantiate_tools(self):
        # This currently is not implemented
        # TODO: complete Tool instantiation
        pass

    def __getattr__(self, name: str) -> Any:
        """
        - Enable dynamic access to agents, tasks or tools as attributes
        - Enable dynamic access to latest information in interactions in the cached session
        """
        try:
            if name.endswith('_agent'):
                return self._agents[name]

            if name.endswith('_task'):
                return self._tasks[name]

            if name.endswith('_tool'):
                return self._tools[name]

            if name.startswith('latest_'):
                element = name[len('latest_'):]
                if self._cached_session.interactions:
                    return self._cached_session.interactions[-1].specialist_results.get(element, '')
                else:
                    return {}

            # Not a known component request
            raise AttributeError(f"'{self.__class__.__name__}' has no attribute '{name}'")
        except KeyError:
            raise AttributeError(f"'{self.__class__.__name__}' has no attribute '{name}'")

    @abstractmethod
    def _instantiate_specialist(self):
        """Instantiate a crew (or flow) to set up the complete specialist, using the assets (agents, tasks, tools).
         The assets can be retrieved using their type name in lower case, e.g. rag_agent"""
        raise NotImplementedError

    def _retrieve_context(self, arguments: SpecialistArguments) -> tuple[str, list[dict[str, Any]]]:
        with current_event.create_span("Specialist Retrieval"):
            self.log_tuning("Starting context retrieval", {
                "num_retrievers": len(self.retrievers),
                "all arguments": arguments.model_dump(),
            })

            modified_arguments = arguments.model_copy(update={
                "query": arguments.question
            })


            # Get retriever-specific arguments
            retriever_arguments = modified_arguments.retriever_arguments

            # Collect context from all retrievers
            all_context = []
            for retriever in self.retrievers:
                # Get arguments for this specific retriever
                retriever_id = str(retriever.retriever_id)
                if retriever_id not in retriever_arguments:
                    current_app.logger.error(f"Missing arguments for retriever {retriever_id}")
                    continue

                # Get the retriever's arguments and update the query
                current_retriever_args = retriever_arguments[retriever_id]
                if isinstance(retriever_arguments[retriever_id], RetrieverArguments):
                    updated_args = current_retriever_args.model_dump()
                    updated_args['question'] = arguments.question
                    updated_args['language'] = arguments.language
                    retriever_args = RetrieverArguments(**updated_args)
                else:
                    # Create a new RetrieverArguments instance from the dictionary
                    current_retriever_args['query'] = arguments.question
                    retriever_args = RetrieverArguments(**current_retriever_args)

                # Each retriever gets its own specific arguments
                retriever_result = retriever.retrieve(retriever_args)
                all_context.extend(retriever_result)

            # Sort by similarity if available and get unique contexts
            all_context.sort(key=lambda x: x.similarity, reverse=True)
            unique_contexts = []
            seen_chunks = set()
            for ctx in all_context:
                if ctx.chunk not in seen_chunks:
                    unique_contexts.append(ctx)
                    seen_chunks.add(ctx.chunk)

            self.log_tuning("Context retrieval completed", {
                "total_contexts": len(all_context),
                "unique_contexts": len(unique_contexts),
                "average_similarity": sum(ctx.similarity for ctx in unique_contexts) / len(
                    unique_contexts) if unique_contexts else 0
            })

        # Prepare context for LLM
        formatted_context = "\n\n".join([
            f"SOURCE: {ctx.metadata.document_id}\n{ctx.chunk}\n\n"
            for ctx in unique_contexts
        ])

        # Return document_ids for citations
        citations = [{"document_id": ctx.metadata.document_id,
                     "document_version_id": ctx.metadata.version_id,
                      "url": ctx.metadata.url}
                     for ctx in unique_contexts]

        self.log_tuning("Context Retrieval Results",
                        {"Formatted Context": formatted_context,
                         "Citations": citations})

        return formatted_context, citations

    def _update_specialist_results(self, specialist_results: SpecialistResult) -> SpecialistResult:
        """Update the specialist results with the latest state information"""
        # Ensure default state result elements are passed
        self._config_default_state_result_relations()

        update_data = {}
        state_dict = self.flow.state.model_dump()
        for state_name, result_name in self._state_result_relations.items():
            if state_name in state_dict and state_dict[state_name] is not None:
                update_data[result_name] = state_dict[state_name]

        return specialist_results.model_copy(update=update_data)

    def _restore_state_from_history(self):
        """Restore the state from the history"""
        if not self._cached_session.interactions:
            return
        last_interaction = self._cached_session.interactions[-1]
        if not last_interaction.specialist_results:
            return
        for state_name, result_name in self._state_result_relations.items():
            if result_name in last_interaction.specialist_results:
                setattr(self.flow.state, state_name, last_interaction.specialist_results[result_name])
        #TODO: Hier wordt steeds een dict of json terug gegeven, geen pydantic model?

        # Initialize the standard state values
        self.flow.state.answer = None
        self.flow.state.question = None
        self.flow.state.form_request = None
        self.flow.state.phase = None
        self.flow.state.citations = []

    @abstractmethod
    def execute(self, arguments: SpecialistArguments,
                formatted_context: Optional[str], citations: Optional[list[dict[str, Any]]]) -> SpecialistResult:
        raise NotImplementedError

    def execute_specialist(self, arguments: SpecialistArguments) -> SpecialistResult:
        self.log_tuning("*****************************************************************************************", {})
        if self.retrievers:
            formatted_context = None
            citations = None
            result = self.execute(arguments, formatted_context, citations)
            modified_result = {
                "citations": citations,
            }
            intermediate_result = result.model_copy(update=modified_result)
        else:
            intermediate_result = self.execute(arguments, "", [])

        final_result = self._update_specialist_results(intermediate_result)

        return final_result
