"""
Built-in system tools for agent task delegation and artifact management.

Provides tools for:
- Delegating tasks to other agents via message queue
- Storing large payloads in artifact registry (MinIO/S3)
- Retrieving artifacts by ID
"""

from __future__ import annotations

from datetime import datetime
import json
from typing import Any
import uuid


class TaskDelegationTool:
    """Tool for delegating tasks to other agents."""

    def __init__(self, message_bus, artifact_storage=None):
        """
        Initialize task delegation tool.
        
        Args:
            message_bus: Message queue backend for inter-agent communication
            artifact_storage: Optional storage backend for large payloads
        """
        self.message_bus = message_bus
        self.artifact_storage = artifact_storage
        self._size_threshold_kb = 100  # Store payloads > 100KB in artifact registry

        # Set tool metadata
        self.delegate_task.__func__.__doc__ = "Delegate a task to another agent. Returns task_id for tracking. Use for complex subtasks or when specialized agents are needed."

    async def delegate_task(
        self,
        agent_name: str,
        task_description: str,
        task_data: dict | None = None,
        wait_for_response: bool = False,
        timeout_seconds: int = 60
    ) -> dict:
        """
        Delegate a task to another agent via message queue.
        
        Args:
            agent_name: Target agent name (e.g., "researcher", "analyzer")
            task_description: Human-readable description of the task
            task_data: Additional structured data for the task
            wait_for_response: If True, wait for agent response synchronously
            timeout_seconds: Max time to wait for response (if wait_for_response=True)
        
        Returns:
            Dict with task_id and optionally response if waiting
        """
        task_data = task_data or {}

        # Build task payload
        payload = {
            "description": task_description,
            "data": task_data,
            "delegated_at": datetime.utcnow().isoformat()
        }

        # Check if payload is too large and should be stored in artifact registry
        payload_size_kb = len(json.dumps(payload).encode()) / 1024
        artifact_id = None

        if self.artifact_storage and payload_size_kb > self._size_threshold_kb:
            # Store in artifact registry
            artifact_id = await self._store_artifact(payload)
            payload = {
                "artifact_id": artifact_id,
                "artifact_size_kb": payload_size_kb,
                "description": task_description
            }

        # Publish task to message queue
        task_id = await self.message_bus.publish_task(agent_name, payload)

        result = {
            "task_id": task_id,
            "agent_name": agent_name,
            "status": "queued"
        }

        if artifact_id:
            result["artifact_id"] = artifact_id
            result["stored_in_registry"] = True

        # Wait for response if requested
        if wait_for_response:
            response = await self.message_bus.wait_for_response(
                task_id,
                timeout_sec=timeout_seconds
            )

            if response:
                result["status"] = "completed"
                result["response"] = response

                # If response contains artifact_id, retrieve it
                if isinstance(response, dict) and "artifact_id" in response:
                    artifact_data = await self._retrieve_artifact(response["artifact_id"])
                    if artifact_data:
                        result["response"] = artifact_data
            else:
                result["status"] = "timeout"
                result["error"] = f"No response after {timeout_seconds}s"

        return result

    async def _store_artifact(self, data: dict) -> str:
        """Store large payload in artifact registry."""
        artifact_id = str(uuid.uuid4())

        if self.artifact_storage:
            await self.artifact_storage.put_object(
                bucket="evenage-artifacts",
                key=f"tasks/{artifact_id}.json",
                data=json.dumps(data).encode()
            )

        return artifact_id

    async def _retrieve_artifact(self, artifact_id: str) -> dict | None:
        """Retrieve artifact from registry."""
        if not self.artifact_storage:
            return None

        try:
            data = await self.artifact_storage.get_object(
                bucket="evenage-artifacts",
                key=f"tasks/{artifact_id}.json"
            )
            return json.loads(data.decode())
        except Exception:
            return None


class ArtifactStorageTool:
    """Tool for storing and retrieving large data artifacts."""

    def __init__(self, storage_backend):
        """
        Initialize artifact storage tool.
        
        Args:
            storage_backend: Storage backend (MinIO/S3)
        """
        self.storage = storage_backend

        # Set tool metadata
        self.store_artifact.__func__.__doc__ = "Store large data (files, results, datasets) in artifact registry. Returns artifact_id for later retrieval. Use for data > 100KB."
        self.retrieve_artifact.__func__.__doc__ = "Retrieve data from artifact registry using artifact_id. Use to fetch previously stored large results."
        self.list_artifacts.__func__.__doc__ = "List available artifacts in the registry. Useful for discovering stored results."

    async def store_artifact(
        self,
        data: str | bytes | dict,
        artifact_type: str = "generic",
        metadata: dict | None = None
    ) -> dict:
        """
        Store large data in artifact registry (MinIO/S3).
        
        Args:
            data: Data to store (string, bytes, or dict which will be JSON-serialized)
            artifact_type: Type label (e.g., "result", "dataset", "file")
            metadata: Optional metadata dict
        
        Returns:
            Dict with artifact_id, size, and storage location
        """
        artifact_id = str(uuid.uuid4())

        # Convert data to bytes
        if isinstance(data, dict):
            data_bytes = json.dumps(data).encode()
        elif isinstance(data, str):
            data_bytes = data.encode()
        else:
            data_bytes = data

        # Store in artifact registry
        bucket = "evenage-artifacts"
        key = f"{artifact_type}/{artifact_id}"

        await self.storage.put_object(
            bucket=bucket,
            key=key,
            data=data_bytes,
            metadata=metadata or {}
        )

        return {
            "artifact_id": artifact_id,
            "type": artifact_type,
            "size_bytes": len(data_bytes),
            "size_kb": round(len(data_bytes) / 1024, 2),
            "location": f"{bucket}/{key}",
            "stored_at": datetime.utcnow().isoformat()
        }

    async def retrieve_artifact(self, artifact_id: str, artifact_type: str = "generic") -> dict:
        """
        Retrieve artifact from storage.
        
        Args:
            artifact_id: Artifact identifier from store_artifact
            artifact_type: Type label used when storing
        
        Returns:
            Dict with artifact data and metadata
        """
        bucket = "evenage-artifacts"
        key = f"{artifact_type}/{artifact_id}"

        try:
            data = await self.storage.get_object(bucket=bucket, key=key)

            # Try to parse as JSON, fall back to string
            try:
                parsed_data = json.loads(data.decode())
            except (json.JSONDecodeError, UnicodeDecodeError):
                parsed_data = data.decode(errors='replace')

            return {
                "artifact_id": artifact_id,
                "type": artifact_type,
                "data": parsed_data,
                "size_bytes": len(data),
                "retrieved_at": datetime.utcnow().isoformat()
            }

        except Exception as e:
            return {
                "artifact_id": artifact_id,
                "error": str(e),
                "status": "not_found"
            }

    async def list_artifacts(
        self,
        artifact_type: str | None = None,
        limit: int = 20
    ) -> dict:
        """
        List artifacts in storage.
        
        Args:
            artifact_type: Optional type filter
            limit: Max number of artifacts to return
        
        Returns:
            Dict with list of artifact metadata
        """
        bucket = "evenage-artifacts"
        prefix = f"{artifact_type}/" if artifact_type else ""

        try:
            objects = await self.storage.list_objects(
                bucket=bucket,
                prefix=prefix,
                max_keys=limit
            )

            artifacts = []
            for obj in objects:
                # Extract artifact_id from key
                parts = obj["key"].split("/")
                artifact_id = parts[-1] if parts else obj["key"]

                artifacts.append({
                    "artifact_id": artifact_id,
                    "type": parts[0] if len(parts) > 1 else "generic",
                    "size_bytes": obj.get("size", 0),
                    "last_modified": obj.get("last_modified")
                })

            return {
                "artifacts": artifacts,
                "count": len(artifacts),
                "bucket": bucket
            }

        except Exception as e:
            return {
                "error": str(e),
                "artifacts": []
            }


def create_system_tools(message_bus, storage_backend=None) -> dict[str, Any]:
    """
    Create system tools for task delegation and artifact management.
    
    Args:
        message_bus: Message queue backend
        storage_backend: Optional storage backend (MinIO/S3)
    
    Returns:
        Dict mapping tool names to tool instances
    """
    tools = {}

    # Task delegation tool
    delegation_tool = TaskDelegationTool(message_bus, storage_backend)
    tools["system_delegate_task"] = delegation_tool.delegate_task

    # Artifact storage tools (if storage backend available)
    if storage_backend:
        artifact_tool = ArtifactStorageTool(storage_backend)
        tools["system_store_artifact"] = artifact_tool.store_artifact
        tools["system_retrieve_artifact"] = artifact_tool.retrieve_artifact
        tools["system_list_artifacts"] = artifact_tool.list_artifacts

    return tools
