# # crew.py file of design_to_code
from crewai import Agent, Crew, Process, Task, LLM
from crewai.project import CrewBase, agent, crew, task
from tools.custom_tool import landing_ai_document_analysis
from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()

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

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

	def __init__(self, anthropic_api_key: str = None):
		"""
		Initialize DocProcessing crew with API key
		
		Args:
			anthropic_api_key: The Anthropic API key to use for LLM
		"""
		# Use provided API key or fall back to environment variable
		api_key = anthropic_api_key or os.getenv("ANTHROPIC_API_KEY")
		
		if not api_key:
			raise ValueError("Anthropic API key is required. Please provide it as parameter or set ANTHROPIC_API_KEY environment variable.")
		
		self.llm = LLM(
			model="claude-3-haiku-20240307",
			api_key=api_key,
			temperature=0,
		)

	@agent
	def document_analyst(self) -> Agent:
		return Agent(
			config=self.agents_config['document_analyst'],
			tools=[landing_ai_document_analysis],
			llm=self.llm,
		)
	
	@agent
	def developer_agent(self) -> Agent:
		return Agent(
			config=self.agents_config['developer_agent'],
			llm=self.llm,
		)

	@task
	def document_analysis(self) -> Task:
		return Task(
			config=self.tasks_config['document_analysis'],
		)
	
	@task
	def code_implementation(self) -> Task:
		return Task(
			config=self.tasks_config['code_implementation'],
		)

	@crew
	def crew(self) -> Crew:
		"""Creates the DocProcessing 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,
		)