import asyncio
import json
import os
import tempfile
import time
from abc import abstractmethod

from crewai import Agent, Task, Crew, Flow
from crewai.agents.parser import AgentAction, AgentFinish
from crewai.crews import CrewOutput
from crewai.tools import BaseTool
from flask import current_app
from pydantic import BaseModel, create_model, Field, ConfigDict, PrivateAttr
from typing import Dict, Type, get_type_hints, Optional, List, Any, Callable

from common.utils.business_event_context import current_event


class EveAICrewAIAgent(Agent):
    specialist: Any = Field(default=None, exclude=True)
    name: str = Field(default=None, exclude=True)
    model_config = ConfigDict(arbitrary_types_allowed=True)

    def __init__(self, specialist, name: str, **kwargs):
        super().__init__(**kwargs)
        self.specialist = specialist
        self.name = name

    def execute_task(
            self,
            task: Task,
            context: Optional[str] = None,
            tools: Optional[List[BaseTool]] = None,
    ) -> str:
        """Execute a task with the agent. Performs AskEveAI specific  fuctionality on top of task execution

        Args:
            task: Task to execute.
            context: Context to execute the task in.
            tools: Tools to use for the task.

        Returns:
            Output of the agent
        """
        # with current_event.create_span(f"Task Execution {task.name} by {self.name}"):
        self.specialist.log_tuning(f"EveAI Agent {self.name}, Task {task.name} Start", {})
        self.specialist.update_progress("EveAI Agent Task Start",
                                        {"name": self.name,
                                         'task': task.name,
                                         })

        result = super().execute_task(task, context, tools)

        self.specialist.log_tuning(f"EveAI Agent {self.name}, Task {task.name} Complete", {})
        self.specialist.update_progress("EveAI Agent Task Complete",
                                        {"name": self.name,
                                         'task': task.name,
                                         })
        return result


class EveAICrewAITask(Task):
    specialist: Any = Field(default=None, exclude=True)
    name: str = Field(default=None, exclude=True)
    model_config = ConfigDict(arbitrary_types_allowed=True)

    def __init__(self, specialist, name: str, **kwargs):
        # kwargs.update({"callback": create_task_callback(self)})
        super().__init__(**kwargs)
        # current_app.logger.debug(f"Task pydantic class for {name}: {"class", self.output_pydantic}")
        self.specialist = specialist
        self.name = name


class EveAICrewAICrew(Crew):
    specialist: Any = Field(default=None, exclude=True)
    name: str = Field(default=None, exclude=True)
    model_config = ConfigDict(arbitrary_types_allowed=True)
    _tmp_dir = PrivateAttr(default=None)

    def __init__(self, specialist, name: str, **kwargs):
        if specialist.tuning:
            base_dir = f"/app/logs/crew_{name}/{specialist.task_id}"
        else:
            self._tmp_dir = tempfile.TemporaryDirectory(prefix="crewai-")
            base_dir = self._tmp_dir.name
        os.environ["CREWAI_STORAGE_DIR"] = base_dir
        os.environ["CREWAI_STORAGE_PATH"] = base_dir
        kwargs['output_log_file'] = f"{base_dir}/specialist_log.json"

        super().__init__(**kwargs)

        self.specialist = specialist
        self.name = name

    async def kickoff_async(
            self,
            inputs: Optional[Dict[str, Any]] = None,
    ) -> CrewOutput:
        async with current_event.create_span_async(f"Crew {self.name} kickoff"):
            start_time = time.time()
            results = await super().kickoff_async(inputs)
            end_time = time.time()
            metrics = {
                "total_tokens": self.usage_metrics.total_tokens,
                "prompt_tokens": self.usage_metrics.prompt_tokens,
                "completion_tokens": self.usage_metrics.completion_tokens,
                "time_elapsed": end_time - start_time,
                "interaction_type": "Crew Execution"
            }
            self.specialist.log_tuning(f"Crew {self.name} async metrics", metrics)
            current_event.log_llm_metrics(metrics)

        return results


class EveAICrewAIFlow(Flow):
    specialist: Any = Field(default=None, exclude=True)
    # name: str = Field(default=None, exclude=True)
    model_config = ConfigDict(arbitrary_types_allowed=True)

    def __init__(self, specialist, name: str, **kwargs):
        super().__init__(**kwargs)
        self.specialist = specialist
        self.name = name
        self.specialist.log_tuning("Initializing EveAICrewAIFlow", {"name": self.name})
        self.specialist.update_progress("EveAI Flow Initialisation", {"name": self.name})

    async def kickoff_async(self, inputs=None) -> Any:
        """Properly override the library's async method"""
        async with current_event.create_span_async(f"Flow {self.name} execution"):
            self.specialist.log_tuning("Inputs retrieved", inputs)
            self.specialist.update_progress("EveAI Flow Start", {"name": self.name})

            try:
                # Call parent's kickoff_async to handle start methods
                result = await super().kickoff_async(inputs)
            except Exception as e:
                current_app.logger.error(f"Error in Flow kickoff_async: {str(e)}")
                raise

            self.specialist.update_progress("EveAI Flow End", {"name": self.name})
            return self.state


class Citation(BaseModel):
    document_id: int
    document_version_id: int
    url: str


class EveAIFlowState(BaseModel):
    """Base class for all EveAI flow states"""
    answer: Optional[str] = None
    question: Optional[str] = None
    phase: Optional[str] = None
    form_request: Optional[Dict[str, Any]] = None
    citations: Optional[List[Citation]] = None

