"""
Genesis Memory System - Base Agent
===================================
Base class for all memory system agents.
Provides common functionality and Claude API integration.
"""

import anthropic
import json
from datetime import datetime, timezone
from typing import Dict, Any, Optional, List
from dataclasses import dataclass, asdict
import os
import sys

sys.path.insert(0, '/mnt/e/genesis-system/genesis-memory')
from core.memory_controller import get_memory_controller


@dataclass
class AgentResult:
    """Result from an agent execution."""
    agent_id: str
    agent_type: str
    team: str
    task: str
    success: bool
    output: str
    tokens_used: int
    cost_usd: float
    duration_ms: int
    memories_created: int
    metadata: Dict[str, Any]


class BaseMemoryAgent:
    """
    Base class for Genesis Memory System agents.

    Provides:
    - Claude API integration
    - Memory controller access
    - Standardized execution pattern
    - Result tracking
    """

    # Pricing (per 1M tokens)
    HAIKU_INPUT_COST = 0.25
    HAIKU_OUTPUT_COST = 1.25
    SONNET_INPUT_COST = 3.0
    SONNET_OUTPUT_COST = 15.0

    def __init__(
        self,
        agent_id: str,
        agent_type: str,
        team: str,
        model: str = "claude-3-5-haiku-20241022",
        system_prompt: Optional[str] = None
    ):
        self.agent_id = agent_id
        self.agent_type = agent_type
        self.team = team
        self.model = model
        self.system_prompt = system_prompt or self._default_system_prompt()

        # Initialize Claude client
        self.client = anthropic.Anthropic()

        # Initialize memory controller
        self.memory = get_memory_controller()

    def _default_system_prompt(self) -> str:
        """Default system prompt for the agent."""
        return f"""You are {self.agent_id}, a specialized agent in the Genesis Memory System.

Team: {self.team}
Role: {self.agent_type}

You have access to a world-class memory system with:
- PostgreSQL for episodic memory
- Redis for working memory
- Qdrant for vector/semantic memory

Your outputs should be clear, actionable, and focused on your specialization.
Always provide structured responses that can be parsed and acted upon."""

    def execute(
        self,
        task: str,
        context: Optional[str] = None,
        max_tokens: int = 2000
    ) -> AgentResult:
        """
        Execute a task using the agent.

        Args:
            task: The task to perform
            context: Additional context to provide
            max_tokens: Maximum response tokens

        Returns:
            AgentResult with execution details
        """
        import time
        start_time = time.time()

        # Build messages
        messages = []

        if context:
            messages.append({
                "role": "user",
                "content": f"Context:\n{context}\n\nTask:\n{task}"
            })
        else:
            messages.append({
                "role": "user",
                "content": task
            })

        try:
            # Call Claude API
            response = self.client.messages.create(
                model=self.model,
                max_tokens=max_tokens,
                system=self.system_prompt,
                messages=messages
            )

            output = response.content[0].text
            input_tokens = response.usage.input_tokens
            output_tokens = response.usage.output_tokens
            total_tokens = input_tokens + output_tokens

            # Calculate cost
            if "haiku" in self.model:
                cost = (input_tokens * self.HAIKU_INPUT_COST +
                       output_tokens * self.HAIKU_OUTPUT_COST) / 1_000_000
            else:
                cost = (input_tokens * self.SONNET_INPUT_COST +
                       output_tokens * self.SONNET_OUTPUT_COST) / 1_000_000

            # Store result in memory
            memory_result = self.memory.remember(
                content=f"Agent {self.agent_id} completed task: {task[:100]}...\nOutput: {output[:500]}...",
                source_type="agent_execution",
                agent_id=self.agent_id,
                importance=0.6
            )

            memories_created = 1 if memory_result.success else 0

            duration_ms = int((time.time() - start_time) * 1000)

            return AgentResult(
                agent_id=self.agent_id,
                agent_type=self.agent_type,
                team=self.team,
                task=task[:200],
                success=True,
                output=output,
                tokens_used=total_tokens,
                cost_usd=cost,
                duration_ms=duration_ms,
                memories_created=memories_created,
                metadata={
                    "model": self.model,
                    "input_tokens": input_tokens,
                    "output_tokens": output_tokens
                }
            )

        except Exception as e:
            duration_ms = int((time.time() - start_time) * 1000)

            return AgentResult(
                agent_id=self.agent_id,
                agent_type=self.agent_type,
                team=self.team,
                task=task[:200],
                success=False,
                output=f"Error: {str(e)}",
                tokens_used=0,
                cost_usd=0.0,
                duration_ms=duration_ms,
                memories_created=0,
                metadata={"error": str(e)}
            )

    def get_memory_context(self, task: str, limit: int = 10) -> str:
        """Get relevant memory context for a task."""
        results = self.memory.recall(
            query=task,
            limit=limit,
            agent_id=self.agent_id
        )

        if not results:
            return "No relevant memories found."

        context_parts = ["Relevant memories:"]
        for r in results:
            context_parts.append(f"- [{r.tier}] {r.content[:200]}...")

        return "\n".join(context_parts)


def run_agent(
    agent_id: str,
    agent_type: str,
    team: str,
    task: str,
    system_prompt: Optional[str] = None
) -> AgentResult:
    """Convenience function to run an agent."""
    agent = BaseMemoryAgent(
        agent_id=agent_id,
        agent_type=agent_type,
        team=team,
        system_prompt=system_prompt
    )
    return agent.execute(task)