# crew.py
from crewai import Agent, Crew, Process, Task, LLM
from crewai.project import CrewBase, agent, crew, task
from dotenv import load_dotenv
import os

from tools.retrieve_financial_data import retrieve_financial_data
from tools.company_profile_tool import get_company_profile
from tools.news_tool import get_stock_news

# Load environment variables
load_dotenv()

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

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

    def __init__(self):
        self.llm = LLM(
            model="claude-3-haiku-20240307",
            api_key=os.getenv("ANTHROPIC_API_KEY")
        )

    @agent
    def stock_information_specialist(self) -> Agent:
        return Agent(
            config=self.agents_config['stock_information_specialist'],
            llm=self.llm,
            tools=[retrieve_financial_data],
            verbose=True
        )

    @agent
    def company_profile_specialist(self) -> Agent:
        return Agent(
            config=self.agents_config['company_profile_specialist'],
            llm=self.llm,
            tools=[get_company_profile],
            verbose=True
        )

    @agent
    def news_analyst(self) -> Agent:
        return Agent(
            config=self.agents_config['news_analyst'],
            llm=self.llm,
            tools=[get_stock_news],
            verbose=True
        )

    @agent
    def investment_advisor(self) -> Agent:
        return Agent(
            config=self.agents_config['investment_advisor'],
            llm=self.llm,
            verbose=True
        )

    @task
    def generate_stock_information(self) -> Task:
        return Task(
            config=self.tasks_config['generate_stock_information'],
            agent=self.stock_information_specialist()
        )

    @task
    def analyze_company_profile(self) -> Task:
        return Task(
            config=self.tasks_config['analyze_company_profile'],
            agent=self.company_profile_specialist()
        )

    @task
    def gather_news_sentiment(self) -> Task:
        return Task(
            config=self.tasks_config['gather_news_sentiment'],
            agent=self.news_analyst()
        )

    @task
    def make_investment_decision(self) -> Task:
        return Task(
            config=self.tasks_config['make_investment_decision'],
            agent=self.investment_advisor(),
            context=[self.generate_stock_information(), self.analyze_company_profile(), self.gather_news_sentiment()]
        )

    @crew
    def crew(self) -> Crew:
        """Creates the PredictingStock crew"""
        return Crew(
            agents=self.agents,  # Automatically created by the @agent decorator
            tasks=self.tasks,    # Automatically created by the @task decorator
            process=Process.sequential,
            verbose=True,
        )