import os
import time
import hashlib
from datetime import datetime
try:
    import anthropic
except ImportError:
    anthropic = None

# RLM Worker (v1.0)
# Uses the highest tier model for specific execution

class RLMWorker:
    def __init__(self, worker_id, specialty="general"):
        self.worker_id = worker_id
        self.specialty = specialty
        # Mapping to the highest tier available via Antigravity Bridge
        self.model = "claude-3-5-sonnet-20241022" 
        self.api_key = os.getenv("ANTHROPIC_API_KEY")
        
        if anthropic and self.api_key:
            self.client = anthropic.Anthropic(api_key=self.api_key)
        else:
            self.client = None

    def execute_task(self, task_description, context=None):
        """Execute task and return output + metadata"""
        start_time = time.time()
        
        if not self.client:
            # Mock behavior if API not available
            output = f"[MOCK] Result of {task_description} from {self.worker_id}"
            usage = type('Usage', (), {'input_tokens': 100, 'output_tokens': 50})
        else:
            # Real API call
            message = f"Context: {context}\n\nTask: {task_description}" if context else task_description
            response = self.client.messages.create(
                model=self.model,
                max_tokens=4096,
                messages=[{"role": "user", "content": message}]
            )
            output = response.content[0].text
            usage = response.usage

        metadata = {
            "worker_id": self.worker_id,
            "specialty": self.specialty,
            "model": self.model,
            "timestamp": datetime.utcnow().isoformat(),
            "execution_time": time.time() - start_time,
            "input_tokens": usage.input_tokens,
            "output_tokens": usage.output_tokens,
            "output_hash": hashlib.sha256(output.encode()).hexdigest(),
            "cost_usd": self._calculate_cost(usage)
        }
        
        return output, metadata

    def _calculate_cost(self, usage):
        # Sonnet 3.5 pricing
        return (usage.input_tokens * 3 / 1_000_000) + (usage.output_tokens * 15 / 1_000_000)