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

Workflow state generators for load testing.
"""

import random
import secrets
import uuid
from datetime import datetime, timedelta
from typing import Generator, List

from sqlalchemy import text

from mcpgateway.db import (
    EmailTeamInvitation,
    EmailTeamJoinRequest,
    TokenRevocation,
    OAuthToken,
)

from ..utils.distributions import exponential_decay_temporal
from .base import BaseGenerator


class TeamInvitationGenerator(BaseGenerator):
    """Generate team invitation records."""

    def get_count(self) -> int:
        """Get total number of team invitations to generate."""
        team_count = self.get_scale_config("teams_per_user_avg", 6) * self.get_scale_config("users", 100)
        invitations_per_team = self.get_scale_config("invitations_per_team_avg", 2)
        return int(team_count * invitations_per_team)

    def get_dependencies(self) -> List[str]:
        """Depends on teams and users."""
        return ["TeamGenerator", "UserGenerator"]

    def generate(self) -> Generator[EmailTeamInvitation, None, None]:
        """Generate team invitation records.

        Yields:
            EmailTeamInvitation instances
        """
        team_result = self.db.execute(text("SELECT id, created_by FROM email_teams ORDER BY created_at"))
        teams = [(row[0], row[1]) for row in team_result.fetchall()]

        if not teams:
            self.logger.warning("No teams found - cannot generate invitations")
            return

        start_date = datetime.fromisoformat(self.config.get("temporal", {}).get("start_date", "2023-01-01"))
        end_date = datetime.fromisoformat(self.config.get("temporal", {}).get("end_date", datetime.now().isoformat()))

        # Generate fake email addresses for invitations
        fake_domains = ["example.com", "test.com", self.email_domain]
        roles = ["viewer", "developer", "admin"]

        invitation_rate = self.get_scale_config("invitation_generation_rate", 0.3)  # 30% of teams have invitations
        teams_with_invitations = random.sample(teams, int(len(teams) * invitation_rate))

        for team_id, created_by in teams_with_invitations:
            num_invitations = random.randint(1, 5)

            for _ in range(num_invitations):
                invited_at = start_date + (end_date - start_date) * random.random()
                expires_at = invited_at + timedelta(days=random.choice([7, 14, 30]))

                # 80% active, 20% expired/used
                is_active = random.random() < 0.8

                yield EmailTeamInvitation(
                    id=str(uuid.uuid4()),
                    team_id=team_id,
                    email=f"invited-{random.randint(1, 100000)}@{random.choice(fake_domains)}",
                    role=random.choice(roles),
                    invited_by=created_by,
                    invited_at=invited_at,
                    expires_at=expires_at,
                    token=secrets.token_urlsafe(32),
                    is_active=is_active,
                )


class TeamJoinRequestGenerator(BaseGenerator):
    """Generate team join request records."""

    def get_count(self) -> int:
        """Get total number of join requests to generate."""
        team_count = self.get_scale_config("teams_per_user_avg", 6) * self.get_scale_config("users", 100)
        requests_per_team = self.get_scale_config("join_requests_per_team_avg", 1)
        return int(team_count * requests_per_team)

    def get_dependencies(self) -> List[str]:
        """Depends on teams and users."""
        return ["TeamGenerator", "UserGenerator"]

    def generate(self) -> Generator[EmailTeamJoinRequest, None, None]:
        """Generate team join request records.

        Yields:
            EmailTeamJoinRequest instances
        """
        team_result = self.db.execute(text("SELECT id FROM email_teams ORDER BY created_at"))
        team_ids = [row[0] for row in team_result.fetchall()]

        user_result = self.db.execute(text("SELECT email FROM email_users ORDER BY created_at"))
        user_emails = [row[0] for row in user_result.fetchall()]

        if not team_ids or not user_emails:
            self.logger.warning("No teams or users found - cannot generate join requests")
            return

        start_date = datetime.fromisoformat(self.config.get("temporal", {}).get("start_date", "2023-01-01"))
        end_date = datetime.fromisoformat(self.config.get("temporal", {}).get("end_date", datetime.now().isoformat()))

        request_rate = self.get_scale_config("join_request_generation_rate", 0.2)  # 20% of teams
        teams_with_requests = random.sample(team_ids, int(len(team_ids) * request_rate))

        statuses = ["pending", "pending", "approved", "rejected"]

        for team_id in teams_with_requests:
            num_requests = random.randint(1, 3)
            selected_users = random.sample(user_emails, min(num_requests, len(user_emails)))

            for user_email in selected_users:
                requested_at = start_date + (end_date - start_date) * random.random()
                expires_at = requested_at + timedelta(days=30)
                status = random.choice(statuses)

                reviewed_at = None
                reviewed_by = None
                if status in ["approved", "rejected"]:
                    reviewed_at = requested_at + timedelta(days=random.randint(1, 7))
                    reviewed_by = random.choice(user_emails)

                yield EmailTeamJoinRequest(
                    id=str(uuid.uuid4()),
                    team_id=team_id,
                    user_email=user_email,
                    message=self.faker.sentence() if random.random() < 0.5 else None,
                    status=status,
                    requested_at=requested_at,
                    expires_at=expires_at,
                    reviewed_at=reviewed_at,
                    reviewed_by=reviewed_by,
                    notes=self.faker.sentence() if reviewed_at and random.random() < 0.3 else None,
                )


class TokenRevocationGenerator(BaseGenerator):
    """Generate token revocation records."""

    def get_count(self) -> int:
        """Get total number of revoked tokens to generate."""
        token_count = self.get_scale_config("users", 100) * self.get_scale_config("tokens_per_user_avg", 3)
        revocation_rate = self.get_scale_config("token_revocation_rate", 0.05)  # 5% revoked
        return int(token_count * revocation_rate)

    def get_dependencies(self) -> List[str]:
        """Depends on tokens."""
        return ["TokenGenerator"]

    def generate(self) -> Generator[TokenRevocation, None, None]:
        """Generate token revocation records.

        Yields:
            TokenRevocation instances
        """
        token_result = self.db.execute(text("SELECT jti, user_email, created_at FROM email_api_tokens ORDER BY created_at"))
        tokens = [(row[0], row[1], datetime.fromisoformat(row[2]) if isinstance(row[2], str) else row[2]) for row in token_result.fetchall()]

        if not tokens:
            self.logger.warning("No tokens found - cannot generate revocations")
            return

        user_result = self.db.execute(text("SELECT email FROM email_users ORDER BY created_at"))
        user_emails = [row[0] for row in user_result.fetchall()]

        revocation_rate = self.get_scale_config("token_revocation_rate", 0.05)
        tokens_to_revoke = random.sample(tokens, int(len(tokens) * revocation_rate))

        for jti, user_email, created_at in tokens_to_revoke:
            # Revoked some time after creation
            revoked_at = created_at + timedelta(days=random.randint(1, 90))

            yield TokenRevocation(
                jti=jti,
                revoked_at=revoked_at,
                revoked_by=user_email,  # Usually self-revoked
                reason=random.choice([
                    "User requested",
                    "Security policy",
                    "Token compromised",
                    "Administrative action",
                    None,
                ]),
            )


class OAuthTokenGenerator(BaseGenerator):
    """Generate OAuth token records for gateway authentication."""

    def get_count(self) -> int:
        """Get total number of OAuth tokens to generate."""
        gateway_count = self.get_scale_config("gateways", 100)
        user_count = self.get_scale_config("users", 100)
        # Not all users have OAuth tokens for all gateways
        oauth_rate = self.get_scale_config("oauth_token_rate", 0.1)  # 10% of user-gateway combinations
        return int(gateway_count * user_count * oauth_rate)

    def get_dependencies(self) -> List[str]:
        """Depends on gateways and users."""
        return ["GatewayGenerator", "UserGenerator"]

    def generate(self) -> Generator[OAuthToken, None, None]:
        """Generate OAuth token records.

        Yields:
            OAuthToken instances
        """
        gateway_result = self.db.execute(text("SELECT id FROM gateways ORDER BY created_at"))
        gateway_ids = [row[0] for row in gateway_result.fetchall()]

        user_result = self.db.execute(text("SELECT email FROM email_users ORDER BY created_at"))
        user_emails = [row[0] for row in user_result.fetchall()]

        if not gateway_ids or not user_emails:
            self.logger.warning("No gateways or users found - cannot generate OAuth tokens")
            return

        start_date = datetime.fromisoformat(self.config.get("temporal", {}).get("start_date", "2023-01-01"))
        end_date = datetime.fromisoformat(self.config.get("temporal", {}).get("end_date", datetime.now().isoformat()))

        oauth_rate = self.get_scale_config("oauth_token_rate", 0.1)

        # Generate random gateway-user pairs
        num_tokens = int(len(gateway_ids) * len(user_emails) * oauth_rate)
        generated_pairs = set()

        for _ in range(num_tokens):
            gateway_id = random.choice(gateway_ids)
            user_email = random.choice(user_emails)
            pair = (gateway_id, user_email)

            if pair in generated_pairs:
                continue

            generated_pairs.add(pair)

            created_at = start_date + (end_date - start_date) * random.random()
            updated_at = created_at + timedelta(days=random.randint(0, 30))
            expires_at = created_at + timedelta(days=random.choice([30, 60, 90, 365]))

            yield OAuthToken(
                id=str(uuid.uuid4()),
                gateway_id=gateway_id,
                user_id=user_email,  # Column is user_id but references email
                app_user_email=user_email,
                access_token=f"ya29.{secrets.token_urlsafe(64)}",
                refresh_token=f"1//{secrets.token_urlsafe(32)}" if random.random() < 0.8 else None,
                token_type="Bearer",
                expires_at=expires_at,
                scopes=["read", "write"] if random.random() < 0.5 else ["read"],
                created_at=created_at,
                updated_at=updated_at,
            )