from crewai import Agent, Crew, Process, Task
from crewai.project import CrewBase, agent, crew, task

from .paths import MARKDOWN_OUTPUT_FILE, PLOT_OUTPUT_FILE


@CrewBase
class Coder():
    """Coder crew"""

    agents_config = 'config/agents.yaml'
    tasks_config = 'config/tasks.yaml'

    # One click install for Docker Desktop:
    #https://docs.docker.com/desktop/

    @agent
    def coder(self) -> Agent:
        from src.inner.tools.custom_tool import (
            ExtractSchema,
            ResultAwareCodeInterpreterTool,
        )

        code_interpreter = ResultAwareCodeInterpreterTool(unsafe_mode=True)
        return Agent(
            config=self.agents_config['coder'],
            verbose=True,
            #allow_code_execution=True,
            #code_execution_mode="unsafe",  # Currently doesn't use Docker for safety
            max_execution_time=300, 
            max_retry_limit=5 ,
            tools=[ExtractSchema(), code_interpreter],
    )


    @task
    def coding_task(self) -> Task:
        task_config = dict(self.tasks_config['coding_task'])
        task_config['description'] = (
            f"{task_config['description']}\n\n"
            f"Save plots to '{PLOT_OUTPUT_FILE}'. "
            f"Write the markdown output to '{MARKDOWN_OUTPUT_FILE}'."
        )
        task_config['output_file'] = str(MARKDOWN_OUTPUT_FILE)

        return Task(
            config=task_config,
        )
    
    # @task
    # def image_task(self) -> Task:
    #     return Task(
    #         config=self.tasks_config['image_task'],
    #     )


    @crew
    def crew(self) -> Crew:
        """Creates the Coder crew"""


        return Crew(
            agents=self.agents, 
            tasks=self.tasks,
            process=Process.sequential,
            verbose=True,
            memory=True,
        )
    
    # TO DO: Change memory to ollama from default (OpenAI)