"""
Usage:
python run_task.py --swagger-file ./configs/sut/openapi.json --scenario-file ./configs/tasks/task_description.json --output-dir output --base-url https://dedicated.staging.tidbapi.com/v1beta1 --sut-token "username:password" --auth-type digest --agent-type llm --llm-model gpt-4o-2024-08-06 --llm-temperature 0.0 --max-steps 20 --region eastus2 --cloud-provider azure --project-id "1369847559694040595"
"""

import argparse
import json
import logging
import sys
from pathlib import Path
from datetime import datetime
from typing import Dict, List

from llmops.task_orchestrator import TaskOrchestrator
from llmops.storage.transition_store import TransitionStore
from dedicated_expert import TidbCloudDedicatedExpert, SutAttr
from llmops.environment.executor import CompoundExecutor


# config logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler("swagger_test_gen.log"),
    ],
)
logger = logging.getLogger(__name__)


def parse_token(token_str: str) -> Dict[str, str]:
    """Parse token string which can be either a JSON string or a regular string"""
    try:
        token_dict = json.loads(token_str)
        if not isinstance(token_dict, dict):
            raise ValueError("Token JSON must be an object {\"role name\": \"token\"}")
        if len(token_dict) == 0:
            raise ValueError("Token JSON object cannot be empty")
        return token_dict
    except json.JSONDecodeError:
        # If not a valid JSON, treat as a regular string and use default role
        return {"org_owner": token_str}


def load_scenarios(file_path: str) -> Dict[str, List[str]]:
    """Load test scenarios"""
    with open(file_path, "r") as f:
        scenarios = json.load(f)
    return {name: scenario["steps"] for name, scenario in scenarios.items()}


def setup_output_dir(base_dir: str) -> Path:
    """set up output directory"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = Path(base_dir) / timestamp
    output_dir.mkdir(parents=True, exist_ok=True)
    return output_dir


def main():
    """main entry"""
    # parse command line arguments
    parser = argparse.ArgumentParser(description="Swagger Test Generator")
    parser.add_argument(
        "--swagger-file",
        required=True,
        help="Path to Swagger/OpenAPI specification file",
    )
    parser.add_argument(
        "--scenario-file", required=True, help="Path to test scenarios definition file"
    )
    parser.add_argument(
        "--output-dir", default="output", help="Directory to store output files"
    )
    parser.add_argument(
        "--base-url", default="http://127.0.0.1:8000", help="Base URL for API server"
    )
    parser.add_argument(
        "--sut-token", default=None, help="Token for API authentication"
    )
    parser.add_argument(
        "--auth-type", default="bearer", help="Type of authentication to use"
    )
    parser.add_argument(
        "--train-episodes", type=int, default=100, help="Number of training episodes"
    )
    parser.add_argument(
        "--max-steps", type=int, default=10, help="Maximum steps per episode"
    )
    parser.add_argument(
        "--agent-type",
        default="epsilon_greedy",
        choices=["epsilon_greedy", "llm"],
        help="Type of agent to use",
    )
    parser.add_argument(
        "--epsilon",
        type=float,
        default=0.1,
        help="Epsilon value for epsilon-greedy agent",
    )
    parser.add_argument(
        "--llm-model", default="gpt-o3-mini", help="LLM model to use for LLM agent"
    )
    parser.add_argument(
        "--llm-temperature",
        type=float,
        default=0.0,
        help="Temperature for LLM sampling",
    )
    parser.add_argument(
        "--region",
        default="us-east-1",
        type=str,
        help="Region for TidbCloud",
    )
    parser.add_argument(
        "--cloud-provider",
        default="aws",
        type=str,
        help="Cloud provider for TidbCloud",
    )
    parser.add_argument(
        "--project-id",
        type=str,
        default="",
        help="Project ID for TidbCloud",
    )
    parser.add_argument(
        "--org-id",
        type=str,
        default="",
        help="Organization ID for TidbCloud",
    )
    args = parser.parse_args()

    try:
        # set up output directory
        output_dir = setup_output_dir(args.output_dir)
        logger.info(f"Output directory: {output_dir}")

        # load test scenarios
        scenarios = load_scenarios(args.scenario_file)
        logger.info(f"Loaded {len(scenarios)} test scenarios")

        # create transition store
        transition_store = TransitionStore(str(output_dir / "transitions.db"))

        # parse token
        token_dict = parse_token(args.sut_token)

        # create business expert
        business_expert = TidbCloudDedicatedExpert(
            args.swagger_file,
            CompoundExecutor(
                base_url=args.base_url,
                api_token=token_dict,
                auth_type=args.auth_type,
            ),
            sut_attr=SutAttr(
                region=args.region,
                cloud_provider=args.cloud_provider,
                project_id=args.project_id,
                org_id=args.org_id,
            ),
        )

        # create task orchestrator
        task_orchestrator = TaskOrchestrator(
            output_dir=output_dir,
            base_url=args.base_url,
            max_steps=args.max_steps,
            agent_type=args.agent_type,
            epsilon=args.epsilon,
            llm_model=args.llm_model,
            llm_temperature=args.llm_temperature,
            transition_store=transition_store,
            business_expert=business_expert,
        )
        task_orchestrator.execute_scenarios(
            scenarios, token_dict, auth_type=args.auth_type
        )

    except Exception as e:
        logger.error(f"Error during execution: {e}", exc_info=True)
        sys.exit(1)


if __name__ == "__main__":
    main()