"""Batch Processing Workflow using Anthropic Batch API.

Enables 50% cost reduction by processing non-urgent tasks asynchronously.
Batch API processes requests within 24 hours - not suitable for interactive workflows.

Copyright 2025 Smart-AI-Memory
Licensed under Fair Source License 0.9
"""

import json
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from attune.config import _validate_file_path
from attune.models import get_model
from attune_llm.providers import AnthropicBatchProvider

logger = logging.getLogger(__name__)


@dataclass
class BatchRequest:
    """Single request in a batch."""

    task_id: str
    task_type: str
    input_data: dict[str, Any]
    model_tier: str = "capable"  # cheap, capable, premium


@dataclass
class BatchResult:
    """Result from batch processing."""

    task_id: str
    success: bool
    output: dict[str, Any] | None = None
    error: str | None = None


class BatchProcessingWorkflow:
    """Process multiple tasks via Anthropic Batch API (50% cost savings).

    Example:
        >>> workflow = BatchProcessingWorkflow()
        >>> requests = [
        ...     BatchRequest(
        ...         task_id="task_1",
        ...         task_type="analyze_logs",
        ...         input_data={"logs": "ERROR: Connection failed..."}
        ...     ),
        ...     BatchRequest(
        ...         task_id="task_2",
        ...         task_type="generate_report",
        ...         input_data={"data": {...}}
        ...     )
        ... ]
        >>> results = await workflow.execute_batch(requests)
        >>> print(f"{sum(r.success for r in results)}/{len(results)} successful")
    """

    def __init__(self, api_key: str | None = None):
        """Initialize batch workflow.

        Args:
            api_key: Anthropic API key (optional, uses ANTHROPIC_API_KEY env var)
        """
        self.batch_provider = AnthropicBatchProvider(api_key=api_key)

    async def execute_batch(
        self,
        requests: list[BatchRequest],
        poll_interval: int = 300,  # 5 minutes
        timeout: int = 86400,  # 24 hours
    ) -> list[BatchResult]:
        """Execute batch of requests.

        Args:
            requests: List of batch requests
            poll_interval: Seconds between status checks (default: 300)
            timeout: Maximum wait time in seconds (default: 86400)

        Returns:
            List of results matching input order

        Raises:
            ValueError: If requests is empty
            TimeoutError: If batch doesn't complete within timeout
            RuntimeError: If batch processing fails

        Example:
            >>> workflow = BatchProcessingWorkflow()
            >>> requests = [
            ...     BatchRequest(
            ...         task_id="task_1",
            ...         task_type="analyze_logs",
            ...         input_data={"logs": "..."}
            ...     )
            ... ]
            >>> results = await workflow.execute_batch(requests)
            >>> for result in results:
            ...     if result.success:
            ...         print(f"Task {result.task_id}: Success")
            ...     else:
            ...         print(f"Task {result.task_id}: {result.error}")
        """
        if not requests:
            raise ValueError("requests cannot be empty")

        # Convert to Anthropic Message Batches format
        api_requests = []
        for req in requests:
            model = get_model("anthropic", req.model_tier)
            if model is None:
                raise ValueError(f"Unknown model tier: {req.model_tier}")

            # Use correct format with params wrapper
            api_requests.append(
                {
                    "custom_id": req.task_id,
                    "params": {
                        "model": model.id,
                        "messages": self._format_messages(req),
                        "max_tokens": 4096,
                    },
                }
            )

        # Submit batch
        logger.info(f"Submitting batch of {len(requests)} requests")
        batch_id = self.batch_provider.create_batch(api_requests)

        logger.info(f"Batch {batch_id} created, polling every {poll_interval}s (max {timeout}s)")

        # Wait for completion
        try:
            raw_results = await self.batch_provider.wait_for_batch(
                batch_id, poll_interval=poll_interval, timeout=timeout
            )
        except TimeoutError:
            logger.error(f"Batch {batch_id} timed out after {timeout}s")
            return [
                BatchResult(
                    task_id=req.task_id,
                    success=False,
                    error="Batch processing timed out",
                )
                for req in requests
            ]
        except RuntimeError as e:
            logger.error(f"Batch {batch_id} failed: {e}")
            return [
                BatchResult(task_id=req.task_id, success=False, error=f"Batch failed: {e}")
                for req in requests
            ]

        # Parse results - new Message Batches API format
        results = []
        for raw in raw_results:
            task_id = raw.get("custom_id", "unknown")
            result = raw.get("result", {})
            result_type = result.get("type", "unknown")

            if result_type == "succeeded":
                # Extract message content from succeeded result
                message = result.get("message", {})
                content_blocks = message.get("content", [])

                # Convert content blocks to simple output format
                output_text = ""
                for block in content_blocks:
                    if isinstance(block, dict) and block.get("type") == "text":
                        output_text += block.get("text", "")

                output = {
                    "content": output_text,
                    "usage": message.get("usage", {}),
                    "model": message.get("model"),
                    "stop_reason": message.get("stop_reason"),
                }
                results.append(BatchResult(task_id=task_id, success=True, output=output))

            elif result_type == "errored":
                # Extract error from errored result
                error = result.get("error", {})
                error_msg = error.get("message", "Unknown error")
                error_type = error.get("type", "unknown_error")
                results.append(
                    BatchResult(task_id=task_id, success=False, error=f"{error_type}: {error_msg}")
                )

            elif result_type == "expired":
                results.append(BatchResult(task_id=task_id, success=False, error="Request expired"))

            elif result_type == "canceled":
                results.append(
                    BatchResult(task_id=task_id, success=False, error="Request canceled")
                )

            else:
                results.append(
                    BatchResult(
                        task_id=task_id,
                        success=False,
                        error=f"Unknown result type: {result_type}",
                    )
                )

        # Log summary
        success_count = sum(r.success for r in results)
        logger.info(f"Batch {batch_id} completed: {success_count}/{len(results)} successful")

        return results

    def _format_messages(self, request: BatchRequest) -> list[dict[str, str]]:
        """Format request into Anthropic messages format.

        Args:
            request: BatchRequest with task_type and input_data

        Returns:
            List of message dicts for Anthropic API
        """
        # Task-specific prompts
        task_prompts = {
            "analyze_logs": "Analyze the following logs and identify issues:\n\n{logs}",
            "generate_report": "Generate a report based on:\n\n{data}",
            "classify_bulk": "Classify the following items:\n\n{items}",
            "generate_docs": "Generate documentation for:\n\n{code}",
            "generate_tests": "Generate unit tests for:\n\n{code}",
            # Add more task types as needed
        }

        # Get prompt template or use default
        prompt_template = task_prompts.get(request.task_type, "Process the following:\n\n{input}")

        # Format with input data
        try:
            content = prompt_template.format(**request.input_data)
        except KeyError as e:
            logger.warning(
                f"Missing required field {e} for task {request.task_type}, using raw input"
            )
            # Use default template instead of the specific one
            default_template = "Process the following:\n\n{input}"
            content = default_template.format(input=json.dumps(request.input_data))

        return [{"role": "user", "content": content}]

    @classmethod
    def from_json_file(cls, file_path: str) -> "BatchProcessingWorkflow":
        """Create workflow from JSON input file.

        Args:
            file_path: Path to JSON file with batch requests

        Returns:
            BatchProcessingWorkflow instance

        Input file format:
            [
                {
                    "task_id": "1",
                    "task_type": "analyze_logs",
                    "input_data": {"logs": "ERROR: ..."},
                    "model_tier": "capable"
                },
                ...
            ]
        """
        return cls()

    def load_requests_from_file(self, file_path: str) -> list[BatchRequest]:
        """Load batch requests from JSON file.

        Args:
            file_path: Path to JSON file

        Returns:
            List of BatchRequest objects

        Raises:
            FileNotFoundError: If file doesn't exist
            json.JSONDecodeError: If file is not valid JSON
            ValueError: If file format is invalid
        """
        path = Path(file_path)
        if not path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        with open(path) as f:
            data = json.load(f)

        if not isinstance(data, list):
            raise ValueError("Input file must contain a JSON array")

        requests = []
        for item in data:
            if not isinstance(item, dict):
                raise ValueError("Each item must be a JSON object")

            requests.append(
                BatchRequest(
                    task_id=item["task_id"],
                    task_type=item["task_type"],
                    input_data=item["input_data"],
                    model_tier=item.get("model_tier", "capable"),
                )
            )

        return requests

    def save_results_to_file(self, results: list[BatchResult], output_path: str) -> None:
        """Save batch results to JSON file.

        Args:
            results: List of BatchResult objects
            output_path: Path to output file

        Raises:
            OSError: If file cannot be written
        """
        output_data = [
            {
                "task_id": r.task_id,
                "success": r.success,
                "output": r.output,
                "error": r.error,
            }
            for r in results
        ]

        path = Path(output_path)
        validated_path = _validate_file_path(str(path))
        with open(validated_path, "w") as f:
            json.dump(output_data, f, indent=2)

        logger.info(f"Results saved to {validated_path}")
