"""
Improved workflow demonstrating rate limiting functionality in Mistral Workflows.
This version includes:
1. Multiple rate-limited activities with different limits (1-5 executions per time window)
2. Workflow that runs based on n_turns parameter
3. Dynamic selection of activity based on rate limit parameter
"""

import asyncio
import json
import time
from uuid import uuid4

import structlog
from pydantic import BaseModel, Field

from mistralai_workflows.core.activity import activity
from mistralai_workflows.core.rate_limiting.rate_limit import RateLimit
from mistralai_workflows.core.workflow import workflow

logger = structlog.get_logger(__name__)


class ActivityParams(BaseModel):
    turn_data: str = Field(..., description="The data for this turn")


# Define multiple rate-limited activities with different limits (1-5)
@activity(rate_limit=RateLimit(time_window_in_sec=1, max_execution=1))
async def process_turn_rate1(params: ActivityParams) -> None:
    """Process turn with rate limit of 1 execution per 2 seconds."""
    pass


@activity(rate_limit=RateLimit(time_window_in_sec=1, max_execution=2))
async def process_turn_rate2(params: ActivityParams) -> None:
    """Process turn with rate limit of 2 executions per 2 seconds."""
    pass


@activity(rate_limit=RateLimit(time_window_in_sec=1, max_execution=3))
async def process_turn_rate3(params: ActivityParams) -> None:
    """Process turn with rate limit of 3 executions per 2 seconds."""
    pass


@activity(rate_limit=RateLimit(time_window_in_sec=1, max_execution=4))
async def process_turn_rate4(params: ActivityParams) -> None:
    """Process turn with rate limit of 4 executions per 2 seconds."""
    pass


@activity(rate_limit=RateLimit(time_window_in_sec=1, max_execution=5))
async def process_turn_rate5(params: ActivityParams) -> None:
    """Process turn with rate limit of 5 executions per 2 seconds."""
    pass


class Params(BaseModel):
    n_turns: int = Field(..., description="Number of turns to process")
    rate_limit: int = Field(1, description="Rate limit to use (1-5 executions per 2 seconds)")


class Result(BaseModel):
    total_turns_processed: int = Field(..., description="Total turns processed")
    processing_time: float = Field(..., description="Total processing time in seconds")
    rate_limit: int = Field(..., description="Rate limit used")


@workflow.define(name="rate-limited-workflow")
class RateLimitedWorkflow:
    """
    A workflow that demonstrates rate limiting functionality with multiple rate limits.
    This workflow:
    1. Processes multiple turns with configurable rate limiting
    2. Dynamically selects the appropriate rate-limited activity
    3. Measures and reports processing metrics
    """

    @workflow.entrypoint
    async def execute(self, params: Params) -> Result:
        """
        Execute the workflow with rate-limited activities.

        Args:
            params.n_turns: Number of turns to process
            params.rate_limit: Rate limit to use (1-5)

        Returns:
            Processing results including metrics
        """
        logger.info("[step 0] Starting workflow", n_turns=params.n_turns, rate_limit=params.rate_limit)

        # Select the appropriate activity based on rate limit parameter
        activity_map = {
            1: process_turn_rate1,
            2: process_turn_rate2,
            3: process_turn_rate3,
            4: process_turn_rate4,
            5: process_turn_rate5,
        }
        selected_activity = activity_map[params.rate_limit]

        logger.info("[step 1] Starting rate-limited workflow execution...")
        start_time = time.time()

        # Process all turns using the selected rate-limited activity
        tasks = []
        for turn in range(params.n_turns):
            turn_data = f"turn-{turn}-rate-{params.rate_limit}"
            tasks.append(selected_activity(ActivityParams(turn_data=turn_data)))

        await asyncio.gather(*tasks)

        duration = time.time() - start_time
        turns_per_second = params.n_turns / duration if duration > 0 else 0

        logger.info(
            "[step 2] Completed processing",
            duration=duration,
            turns_per_second=turns_per_second,
            expected_rate=params.rate_limit / 2,  # Since time window is 2 seconds
        )

        return Result(total_turns_processed=params.n_turns, processing_time=duration, rate_limit=params.rate_limit)


async def main() -> None:
    """Run the rate-limited workflow example."""
    import os

    os.environ["DANGEROUSLY_FORCE_FAIL_WORKFLOW_ON_ERROR"] = "true"
    os.environ["RETRY_POLICY_MAX_ATTEMPTS"] = "1"

    from mistralai_workflows.core.config.config import config
    from mistralai_workflows.core.temporal.temporal_client import create_temporal_client

    start_time = time.time()
    client = await create_temporal_client()

    # Start the worker
    from mistralai_workflows.core.worker import run_worker

    worker_task = await run_worker([RateLimitedWorkflow], detach=True)
    assert worker_task is not None

    try:
        # Execute the workflow with different rate limits
        n_turns = 100
        logger.info(f"Starting workflows with {n_turns} turns each")
        results = []
        for _ in range(3):
            results_ = await asyncio.gather(
                *[
                    client.execute_workflow(
                        RateLimitedWorkflow.execute,
                        Params(n_turns=n_turns, rate_limit=rate_limit),
                        id=f"rate-limited-workflow-{uuid4()}",
                        task_queue=config.temporal.namespace,
                    )
                    for rate_limit in [1, 2, 3, 4, 5]
                ]
            )
            results.extend(results_)
        json_result = [r.model_dump() for r in results]
        print(json.dumps(json_result))  # noqa: T201
    finally:
        # Clean up
        worker_task.cancel()
        try:
            await worker_task
        except asyncio.CancelledError:
            pass

    duration = round(time.time() - start_time)
    logger.info(f"All workflows completed in {duration} seconds")


if __name__ == "__main__":
    asyncio.run(main())
