"""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 the Apache License, Version 2.0
"""

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

from attune.llm.providers import AnthropicBatchProvider
from attune.models import get_model
from attune.security.path_validation import _validate_file_path

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.exception(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.exception(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}")
