from crewai import Agent, Crew, Process, Task, LLM
from crewai.project import CrewBase, agent, crew, task
from crewai.agents.agent_builder.base_agent import BaseAgent
from crewai_tools import FileReadTool
from typing import List
import os

# If you want to run a snippet of code before or after the crew starts,
# you can use the @before_kickoff and @after_kickoff decorators
# https://docs.crewai.com/concepts/crews#example-crew-class-with-decorators

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

    agents: List[BaseAgent]
    tasks: List[Task]
    
    llm = LLM(
        model=os.getenv('MODEL'), 
        aws_region_name="us-east-2",
        aws_profile_name=os.getenv("AWS_PROFILE", "default"),  # Uses AWS CLI profile
        temperature=0.2,
        top_p=None
    )


    # Learn more about YAML configuration files here:
    # Agents: https://docs.crewai.com/concepts/agents#yaml-configuration-recommended
    # Tasks: https://docs.crewai.com/concepts/tasks#yaml-configuration-recommended
    
    # If you would like to add tools to your agents, you can learn more about it here:
    # https://docs.crewai.com/concepts/agents#agent-tools
    @agent
    def researcher(self) -> Agent:
        return Agent(
            config=self.agents_config['researcher'], # type: ignore[index]
            tools=[FileReadTool()],
            llm=self.llm,
            verbose=True
        )

    @agent
    def analyzer(self) -> Agent:
        return Agent(
            config=self.agents_config['analyzer'], # type: ignore[index]
            llm=self.llm,
            verbose=True
        )

    @agent
    def reporting_analyst(self) -> Agent:
        return Agent(
            config=self.agents_config['reporting_analyst'], # type: ignore[index]
            llm=self.llm,
            verbose=True
        )

    # To learn more about structured task outputs,
    # task dependencies, and task callbacks, check out the documentation:
    # https://docs.crewai.com/concepts/tasks#overview-of-a-task
    @task
    def research_task(self) -> Task:
        return Task(
            config=self.tasks_config['research_task'], # type: ignore[index]
        )

    @task
    def analysis_task(self) -> Task:
        return Task(
            config=self.tasks_config['analysis_task'], # type: ignore[index]
        )

    @task
    def reporting_task(self) -> Task:
        return Task(
            config=self.tasks_config['reporting_task'], # type: ignore[index]
            output_file='report.md'
        )

    @crew
    def crew(self) -> Crew:
        """Creates the ValkeyResearchDemo crew with LanceDB storage (testing without Valkey)"""
        from crewai.memory.storage.lancedb_storage import LanceDBStorage
        from crewai.memory import Memory
        
        # Use LanceDB storage (same backend as Valkey) with Bedrock embeddings
        lancedb_storage = LanceDBStorage(
            path="./lancedb_test"
        )
        
        # Wrap storage in Memory with embedder config
        memory = Memory(
            storage=lancedb_storage,
            llm=self.llm,
            embedder={
                "provider": "amazon-bedrock",
                "config": {"model": "amazon.titan-embed-text-v2:0"}
            }
        )
        
        print("✓ Using LanceDB storage with Bedrock embeddings (testing core crewai main branch)")
        
        return Crew(
            agents=self.agents, # Automatically created by the @agent decorator
            tasks=self.tasks, # Automatically created by the @task decorator
            process=Process.sequential,
            memory=memory,  # Use Memory with LanceDB storage
            llm=self.llm,
            verbose=True,
            # process=Process.hierarchical, # In case you wanna use that instead https://docs.crewai.com/how-to/Hierarchical/
        )

---

researcher:
  role: >
    {topic} Senior Data Researcher
  goal: >
    Uncover cutting-edge developments in {topic}
  backstory: >
    You're a seasoned researcher with a knack for uncovering the latest
    developments in {topic}. Known for your ability to find the most relevant
    information and present it in a clear and concise manner.

analyzer:
  role: >
    {topic} Data Analyzer
  goal: >
    Analyze research findings and identify key patterns, trends, and insights about {topic}
  backstory: >
    You're an expert data analyst who excels at synthesizing complex information.
    You can quickly identify patterns, validate findings, and determine which
    insights are most valuable for decision-making.

reporting_analyst:
  role: >
    {topic} Reporting Analyst
  goal: >
    Create detailed reports based on {topic} data analysis and research findings
  backstory: >
    You're a meticulous analyst with a keen eye for detail. You're known for
    your ability to turn complex data into clear and concise reports, making
    it easy for others to understand and act on the information you provide.

---

researcher:
  role: >
    {topic} Senior Data Researcher
  goal: >
    Uncover cutting-edge developments in {topic}
  backstory: >
    You're a seasoned researcher with a knack for uncovering the latest
    developments in {topic}. Known for your ability to find the most relevant
    information and present it in a clear and concise manner.

analyzer:
  role: >
    {topic} Data Analyzer
  goal: >
    Analyze research findings and identify key patterns, trends, and insights about {topic}
  backstory: >
    You're an expert data analyst who excels at synthesizing complex information.
    You can quickly identify patterns, validate findings, and determine which
    insights are most valuable for decision-making.

reporting_analyst:
  role: >
    {topic} Reporting Analyst
  goal: >
    Create detailed reports based on {topic} data analysis and research findings
  backstory: >
    You're a meticulous analyst with a keen eye for detail. You're known for
    your ability to turn complex data into clear and concise reports, making
    it easy for others to understand and act on the information you provide.