# -*- coding: utf-8 -*-
"""Location: ./tests/load/generate.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

Main CLI for generating load test data.

Usage:
    python -m tests.load.generate --profile production
    python -m tests.load.generate --config tests/load/configs/custom.yaml
"""

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

import yaml
from faker import Faker
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from mcpgateway.config import settings

from .generators import (
    UserGenerator,
    TeamGenerator,
    TeamMemberGenerator,
    TokenGenerator,
    GatewayGenerator,
    ToolGenerator,
    ResourceGenerator,
    PromptGenerator,
    ServerGenerator,
    A2AAgentGenerator,
    ServerToolAssociationGenerator,
    ServerResourceAssociationGenerator,
    ServerPromptAssociationGenerator,
    ServerA2AAssociationGenerator,
    ToolMetricsGenerator,
    ResourceMetricsGenerator,
    PromptMetricsGenerator,
    ServerMetricsGenerator,
    A2AAgentMetricsGenerator,
    TokenUsageLogGenerator,
    EmailAuthEventGenerator,
    PermissionAuditLogGenerator,
    MCPSessionGenerator,
    MCPMessageGenerator,
    ResourceSubscriptionGenerator,
    TeamInvitationGenerator,
    TeamJoinRequestGenerator,
    TokenRevocationGenerator,
    OAuthTokenGenerator,
)
from .utils.progress import MultiProgressTracker, SimpleProgressTracker
from .utils.validation import DataValidator
from rich.console import Console


# Logger setup
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


# Generator registry
GENERATORS = {
    # Core entities
    "users": UserGenerator,
    "teams": TeamGenerator,
    "team_members": TeamMemberGenerator,
    "tokens": TokenGenerator,
    "gateways": GatewayGenerator,
    "tools": ToolGenerator,
    "resources": ResourceGenerator,
    "prompts": PromptGenerator,
    "servers": ServerGenerator,
    "a2a_agents": A2AAgentGenerator,
    # Associations
    "server_tool_associations": ServerToolAssociationGenerator,
    "server_resource_associations": ServerResourceAssociationGenerator,
    "server_prompt_associations": ServerPromptAssociationGenerator,
    "server_a2a_associations": ServerA2AAssociationGenerator,
    # Metrics
    "tool_metrics": ToolMetricsGenerator,
    "resource_metrics": ResourceMetricsGenerator,
    "prompt_metrics": PromptMetricsGenerator,
    "server_metrics": ServerMetricsGenerator,
    "a2a_agent_metrics": A2AAgentMetricsGenerator,
    # Activity logs
    "token_usage_logs": TokenUsageLogGenerator,
    "email_auth_events": EmailAuthEventGenerator,
    "permission_audit_logs": PermissionAuditLogGenerator,
    # Sessions
    "mcp_sessions": MCPSessionGenerator,
    "mcp_messages": MCPMessageGenerator,
    "resource_subscriptions": ResourceSubscriptionGenerator,
    # Workflow state
    "team_invitations": TeamInvitationGenerator,
    "team_join_requests": TeamJoinRequestGenerator,
    "token_revocations": TokenRevocationGenerator,
    "oauth_tokens": OAuthTokenGenerator,
}


def load_config(config_path: str) -> Dict[str, Any]:
    """Load configuration from YAML file.

    Args:
        config_path: Path to YAML config file

    Returns:
        Configuration dictionary
    """
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)


def get_profile_path(profile: str) -> str:
    """Get path to profile configuration.

    Args:
        profile: Profile name (small, medium, large, production)

    Returns:
        Path to profile config file
    """
    base_path = Path(__file__).parent / "configs"
    return str(base_path / f"{profile}.yaml")


def create_database_session(config: Dict[str, Any]):
    """Create database session from configuration.

    Args:
        config: Configuration dictionary

    Returns:
        SQLAlchemy session
    """
    database_url = settings.database_url
    pool_size = config.get("database", {}).get("pool_size", 20)
    max_overflow = config.get("database", {}).get("max_overflow", 40)
    pool_timeout = config.get("database", {}).get("pool_timeout", 30)

    engine = create_engine(
        database_url,
        pool_size=pool_size,
        max_overflow=max_overflow,
        pool_timeout=pool_timeout,
        echo=False,
    )

    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
    return SessionLocal()


def generate_data(config: Dict[str, Any], dry_run: bool = False) -> Dict[str, Any]:
    """Generate load test data.

    Args:
        config: Configuration dictionary
        dry_run: If True, don't actually insert data

    Returns:
        Generation statistics
    """
    # Initialize console and progress tracking
    console = Console()
    simple_progress = SimpleProgressTracker(console)

    profile_name = config.get('profile', {}).get('name', 'unknown')
    simple_progress.print_step(f"Starting load data generation (Profile: {profile_name})")

    logger.info("Starting load data generation...")
    logger.info(f"Profile: {profile_name}")

    start_time = time.time()
    stats = {}

    # Setup
    db = create_database_session(config)
    faker = Faker(config.get("realism", {}).get("faker_locale", "en_US"))

    # Set random seed for reproducibility
    seed = config.get("global", {}).get("random_seed", 42)
    Faker.seed(seed)

    # Determine generation order (respects dependencies)
    generation_order = config.get("generation_order", list(GENERATORS.keys()))

    # Initialize multi-progress tracker if not dry run
    progress_tracker = None
    use_progress = config.get("global", {}).get("progress_bar", True) and not dry_run

    if use_progress:
        progress_tracker = MultiProgressTracker(console)

        # Pre-calculate all generator counts and add tasks
        simple_progress.print_step("Calculating record counts...")
        for generator_name in generation_order:
            if generator_name in GENERATORS:
                generator_class = GENERATORS[generator_name]
                temp_generator = generator_class(db, config, faker, logger)
                count = temp_generator.get_count()
                progress_tracker.add_task(
                    generator_name,
                    count,
                    f"{generator_name:30}"
                )

        simple_progress.print_success(f"Planning to generate {progress_tracker.total_records:,} total records")

    try:
        if use_progress and progress_tracker:
            # Use rich live display for progress
            with progress_tracker.live_display():
                for generator_name in generation_order:
                    if generator_name not in GENERATORS:
                        logger.warning(f"Unknown generator: {generator_name}")
                        continue

                    generator_class = GENERATORS[generator_name]
                    generator = generator_class(db, config, faker, logger, progress_tracker=progress_tracker)

                    # Run generator with progress tracking
                    gen_stats = generator.run()
                    stats[generator_name] = gen_stats

                    # Explicitly refresh display after each generator completes
                    progress_tracker.refresh()
        else:
            # Dry run or no progress bar - simple output
            for generator_name in generation_order:
                if generator_name not in GENERATORS:
                    logger.warning(f"Unknown generator: {generator_name}")
                    continue

                generator_class = GENERATORS[generator_name]
                generator = generator_class(db, config, faker, logger)
                count = generator.get_count()

                if dry_run:
                    simple_progress.print_info(f"[DRY RUN] Would generate {count:,} {generator_name}")
                    stats[generator_name] = {"generated": count, "inserted": 0, "dry_run": True, "duration": 0}
                    continue

                simple_progress.print_step(f"Generating {generator_name}...")
                gen_stats = generator.run()
                stats[generator_name] = gen_stats
                simple_progress.print_success(f"Completed {generator_name}: {gen_stats['generated']:,} records")

        # Calculate totals
        total_generated = sum(s.get("generated", 0) for s in stats.values())
        total_inserted = sum(s.get("inserted", 0) for s in stats.values())

        elapsed_time = time.time() - start_time

        summary = {
            "profile": config.get("profile", {}).get("name", "custom"),
            "timestamp": datetime.now().isoformat(),
            "duration_seconds": elapsed_time,
            "total_generated": total_generated,
            "total_inserted": total_inserted,
            "records_per_second": total_generated / elapsed_time if elapsed_time > 0 else 0,
            "generator_stats": stats,
            "dry_run": dry_run,
        }

        # Print summary
        console.print()
        simple_progress.print_summary(total_generated, elapsed_time, profile_name)

        # Validate if enabled
        if config.get("validation", {}).get("enabled", True) and not dry_run:
            console.print()
            simple_progress.print_step("Running validation checks...")
            validator = DataValidator(db, logger)
            validation_results = validator.validate_all()
            summary["validation"] = validation_results

            if validation_results.get("all_passed", False):
                simple_progress.print_success("All validation checks passed!")
            else:
                simple_progress.print_warning("Some validation checks failed")

        return summary

    except Exception as e:
        logger.error(f"Generation failed: {e}", exc_info=True)
        simple_progress.print_error(f"Generation failed: {e}")
        db.rollback()
        raise
    finally:
        db.close()


def main():
    """Main CLI entry point."""
    parser = argparse.ArgumentParser(
        description="Generate production-scale load test data for ContextForge",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    parser.add_argument(
        "--profile",
        type=str,
        choices=["small", "medium", "large", "production", "massive"],
        help="Load profile to use"
    )

    parser.add_argument(
        "--config",
        type=str,
        help="Path to custom configuration file"
    )

    parser.add_argument(
        "--workers",
        type=int,
        default=1,
        help="Number of parallel workers (default: 1)"
    )

    parser.add_argument(
        "--batch-size",
        type=int,
        help="Batch insert size (overrides config)"
    )

    parser.add_argument(
        "--seed",
        type=int,
        help="Random seed for reproducibility"
    )

    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show what would be generated without inserting"
    )

    parser.add_argument(
        "--skip-validation",
        action="store_true",
        help="Skip post-generation validation"
    )

    parser.add_argument(
        "--output",
        type=str,
        help="Output report path (JSON)"
    )

    parser.add_argument(
        "--log-level",
        type=str,
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        default="INFO",
        help="Logging level"
    )

    args = parser.parse_args()

    # Set log level
    logging.getLogger().setLevel(getattr(logging, args.log_level))

    # Load configuration
    if args.config:
        config_path = args.config
    elif args.profile:
        config_path = get_profile_path(args.profile)
    else:
        logger.error("Must specify either --profile or --config")
        sys.exit(1)

    if not Path(config_path).exists():
        logger.error(f"Configuration file not found: {config_path}")
        sys.exit(1)

    config = load_config(config_path)

    # Apply CLI overrides
    if args.batch_size:
        config.setdefault("global", {})["batch_size"] = args.batch_size

    if args.seed is not None:
        config.setdefault("global", {})["random_seed"] = args.seed

    if args.skip_validation:
        config.setdefault("validation", {})["enabled"] = False

    # Generate data
    try:
        summary = generate_data(config, dry_run=args.dry_run)

        # Save report
        output_path = args.output or config.get("reporting", {}).get("output_file")
        if output_path:
            output_file = Path(output_path)
            output_file.parent.mkdir(parents=True, exist_ok=True)

            with open(output_file, 'w') as f:
                json.dump(summary, f, indent=2, default=str)

            logger.info(f"Report saved to: {output_file}")

        sys.exit(0)

    except Exception as e:
        logger.error(f"Failed: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()