"""
Environment Variables Configuration:
--------------------------------------------------------------
1. DASHSCOPE_API_KEY (Required)
   - Sets the authentication key for DashScope API access
   - Format: API key string (e.g., "ds-1234567890abcdef")

2. LOG_LEVEL (Optional, default: 'info')
   - Controls logging verbosity
   - Allowed values: 'debug', 'info', 'warning', 'error', 'critical'

3. FC_PYPI_LIB (Required)
   - Specifies local/alternative dashscope package installations
   - Accepts space-separated list of wheel paths
   - Default: Uses official PyPI 'dashscope' package
   - Example: "dashscope-1.25.16-py3-none-any.whl"

Functions Workflow (main_functions):
--------------------------------------------------------------
register_functions():
- Registers and initializes function components for RL workflow
- Parameters:
  - functions: List of AgenticRLFunctionComponent configurations
  - lazy_load: If False, immediately loads components (default: False)
- Returns tuple of registered component IDs (rollout/reward/group_reward entities and instances)

test_functions():
- Validates registered function components
- Parameters:
  - instance_id: Target component instance ID
  - type: FunctionType (ROLLOUT/REWARD/GROUP_REWARD)
  - input_data: Test input payload
- Verifies component outputs against expected results
"""

import json
from typing import List

from dashscope.finetune.reinforcement import (logger, TaskStatus, FunctionType, Datasets, DatasetsType, DataSourceType, Dataset, MountStorage,
                                              AgenticRLTuning, AgenticRLFunctionComponent, FunctionComponentModel, FunctionComponentRuntime,
                                              RolloutInput, RewardInput, RewardOutput, RolloutOutput)
from dashscope.finetune.agentic_rl import AgenticRL
from dashscope.finetune.finetunes import FineTunes


async def main_functions():
    """Main execution workflow"""
    try:
        logger.info("Starting main tests functions")
        client = AgenticRL()

        functions = [
            # rollout-only
            AgenticRLFunctionComponent(
                type=FunctionType.ROLLOUT,
                fcmodel=FunctionComponentModel(
                    classpath="functions.rollout.rollout_only.DemoRolloutProcessor"),
            ),
            # reward
            AgenticRLFunctionComponent(
                type=FunctionType.REWARD,
                fcmodel=FunctionComponentModel(
                    classpath="functions.reward.reward.DemoRewardProcessor"),
            ),
            # reward-decorator
            AgenticRLFunctionComponent(
                type=FunctionType.REWARD,
                fcmodel=FunctionComponentModel(
                    classpath="functions/reward/reward_decorator.py:SafetyProcessor"),
            ),
            # group-reward
            AgenticRLFunctionComponent(
                type=FunctionType.GROUP_REWARD,
                fcmodel=FunctionComponentModel(
                    classpath="functions.reward.group_reward.DemoGroupRewardProcessor"),
            ),
        ]

        # Register functions
        (rollout_entity_ids,
         reward_entity_ids,
         group_reward_entity_ids,
         rollout_instance_ids,
         reward_instance_ids,
         group_reward_instance_ids) = \
            await client.register_functions(
                functions=functions,
                lazy_load=False) # register & load functions

        logger.info(
            f"agentic rl register functions: {rollout_entity_ids=}, {reward_entity_ids=}, {group_reward_entity_ids=}, {rollout_instance_ids=}, {reward_instance_ids=}, {group_reward_instance_ids=}")

        with open("./resources/rollout_input.json", "r") as file:
            json_data = json.load(file)
            rollout_input = json_data

        with open("./resources/reward_input.json", "r") as file:
            json_data = json.load(file)
            reward_input = json_data

        with open("./resources/reward_decorator_input.json", "r") as file:
            json_data = json.load(file)
            reward_decorator_input = json_data

        with open("./resources/group_reward_input.json", "r") as file:
            json_data = json.load(file)
            group_reward_input = json_data

        # Testing rollout-only functions
        if rollout_instance_ids and rollout_instance_ids[0]:
            result = await AgenticRL.test_functions(instance_id=rollout_instance_ids[0],
                                                    type=FunctionType.ROLLOUT,
                                                    input_data=rollout_input)
            logger.info(f"agentic rl test rollout: {rollout_instance_ids[0]=}, {result=}")
            status = result.get('status', None)
            assert status == TaskStatus.SUCCESS, f"Expected status, got {status}"

        # Testing reward functions
        if reward_instance_ids and reward_instance_ids[0]:
            result = await AgenticRL.test_functions(instance_id=reward_instance_ids[0],
                                                    type=FunctionType.REWARD,
                                                    input_data=reward_input)
            logger.info(f"agentic rl test rewards: {reward_instance_ids[0]=}, {result=}")
            status = result.get('status', None)
            assert status == TaskStatus.SUCCESS, f"Expected status, got {status}"

        # Testing reward-decorator functions
        if reward_instance_ids and reward_instance_ids[1]:
            result = await AgenticRL.test_functions(instance_id=reward_instance_ids[1],
                                                    type=FunctionType.REWARD,
                                                    input_data=reward_decorator_input)
            logger.info(f"agentic rl test rewards-decorator: {reward_instance_ids[1]=}, {result=}")
            reward_score = result.get('reward', {}).get('reward_score', 0.0)
            assert reward_score == 0.85, f"Expected reward_score 0.85, got {reward_score}"

        # Testing group-reward functions
        if group_reward_instance_ids and group_reward_instance_ids[0]:
            result = await AgenticRL.test_functions(instance_id=group_reward_instance_ids[0],
                                                    type=FunctionType.GROUP_REWARD,
                                                    input_data=group_reward_input)
            logger.info(f"agentic rl test group-rewards: {group_reward_instance_ids[0]=}, {result=}")
            status = result.get('status', None)
            assert status == TaskStatus.SUCCESS, f"Expected status, got {status}"

        logger.info("All tests functions completed successfully")

    except Exception as e:
        logger.error(f"Main execution flow terminated: {e}")


if __name__ == "__main__":
    import asyncio

    asyncio.run(main_functions())
