#!/usr/bin/env python3
"""
Generate OAuth2 access tokens for MCP agents using Keycloak
Python version of generate-agent-token.sh with batch processing capabilities
"""

import argparse
import glob
import json
import logging
import os
import sys
from datetime import UTC, datetime
from typing import Any

import requests


class Colors:
    """ANSI color codes for console output"""

    RED = "\033[0;31m"
    GREEN = "\033[0;32m"
    YELLOW = "\033[1;33m"
    BLUE = "\033[0;34m"
    NC = "\033[0m"  # No Color


class TokenGenerator:
    """Generate tokens for MCP agents using Keycloak OAuth2"""

    def __init__(self, verbose: bool = False):
        self.verbose = verbose
        self.setup_logging()

    def setup_logging(self):
        """Setup logging configuration"""
        level = logging.DEBUG if self.verbose else logging.INFO
        logging.basicConfig(level=level, format="%(asctime)s - %(levelname)s - %(message)s")
        self.logger = logging.getLogger(__name__)

    def log(self, message: str):
        """Log info message if verbose mode is enabled"""
        if self.verbose:
            self.logger.debug(message)

    def error(self, message: str):
        """Print error message"""
        self.logger.error(message)

    def success(self, message: str):
        """Print success message"""
        self.logger.info(message)

    def warning(self, message: str):
        """Print warning message"""
        self.logger.warning(message)

    def load_agent_config(self, agent_name: str, oauth_tokens_dir: str) -> dict[str, Any] | None:
        """Load agent configuration from JSON file"""
        config_file = os.path.join(oauth_tokens_dir, f"{agent_name}.json")

        if not os.path.exists(config_file):
            self.error(f"Config file not found: {config_file}")
            return None

        self.log(f"Loading config from: {config_file}")

        try:
            with open(config_file) as f:
                config = json.load(f)
            return config
        except json.JSONDecodeError as e:
            self.error(f"Failed to parse JSON config file: {e}")
            return None
        except Exception as e:
            self.error(f"Failed to load config file: {e}")
            return None

    def get_token_from_keycloak(
        self, client_id: str, client_secret: str, keycloak_url: str, realm: str
    ) -> dict[str, Any] | None:
        """Request access token from Keycloak"""
        token_url = f"{keycloak_url}/realms/{realm}/protocol/openid-connect/token"

        self.log(f"Token URL: {token_url}")
        self.log(f"Client ID: {client_id}")
        self.log(f"Realm: {realm}")

        data = {
            "grant_type": "client_credentials",
            "client_id": client_id,
            "client_secret": client_secret,
            "scope": "openid email profile",
        }

        headers = {"Content-Type": "application/x-www-form-urlencoded"}

        try:
            response = requests.post(token_url, data=data, headers=headers, timeout=30)
            response.raise_for_status()

            token_data = response.json()

            # Check for error in response
            if "error_description" in token_data:
                self.error(f"Token request failed: {token_data['error_description']}")
                return None

            # Validate access token exists
            if "access_token" not in token_data:
                self.error("No access token in response")
                self.log(f"Response keys: {list(token_data.keys())}")
                return None

            return token_data

        except requests.exceptions.RequestException as e:
            self.error(f"Failed to make token request to Keycloak: {e}")
            return None
        except json.JSONDecodeError as e:
            self.error(f"Invalid JSON response: {e}")
            return None

    def save_token_files(
        self,
        agent_name: str,
        token_data: dict[str, Any],
        client_id: str,
        client_secret: str,
        keycloak_url: str,
        realm: str,
        oauth_tokens_dir: str,
    ) -> bool:
        """Save token to both .env and .json files"""
        access_token = token_data["access_token"]
        expires_in = token_data.get("expires_in")

        # Create output directory
        os.makedirs(oauth_tokens_dir, exist_ok=True)

        # Generate timestamps
        generated_at = datetime.now(UTC).isoformat()
        expires_at = None
        if expires_in:
            expiry_timestamp = datetime.now(UTC).timestamp() + expires_in
            expires_at = datetime.fromtimestamp(expiry_timestamp, UTC).isoformat()

        # Save .env file with restricted permissions (contains secrets)
        env_file = os.path.join(oauth_tokens_dir, f"{agent_name}.env")
        try:
            with open(env_file, "w") as f:  # nosec - intentional credential storage for CLI token cache
                f.write(f"# Generated access token for {agent_name}\n")
                f.write(f"# Generated at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f'export ACCESS_TOKEN="{access_token}"\n')  # nosec - intentional token storage in secured file
                f.write(f'export CLIENT_ID="{client_id}"\n')
                f.write(f'export CLIENT_SECRET="{client_secret}"\n')  # nosec - intentional credential storage in secured file
                f.write(f'export KEYCLOAK_URL="{keycloak_url}"\n')
                f.write(f'export KEYCLOAK_REALM="{realm}"\n')
                f.write('export AUTH_PROVIDER="keycloak"\n')
            os.chmod(env_file, 0o600)  # Restrict file permissions to owner only
        except Exception as e:
            self.error(f"Failed to save .env file: {e}")
            return False

        # Save .json file with metadata
        json_file = os.path.join(oauth_tokens_dir, f"{agent_name}-token.json")
        token_json = {
            "agent_name": agent_name,
            "access_token": access_token,
            "token_type": "Bearer",  # nosec B105 - OAuth2 standard token type per RFC 6750
            "expires_in": expires_in,
            "generated_at": generated_at,
            "expires_at": expires_at,
            "provider": "keycloak",
            "keycloak_url": keycloak_url,
            "keycloak_realm": realm,
            "client_id": client_id,
            "scope": "openid email profile",
            "metadata": {
                "generated_by": "generate_tokens.py",
                "script_version": "1.0",
                "token_format": "JWT",  # nosec B105 - Standard token format identifier, not a password
                "auth_method": "client_credentials",
            },
        }

        try:
            with open(json_file, "w") as f:
                json.dump(token_json, f, indent=2)
            os.chmod(json_file, 0o600)  # Restrict file permissions to owner only
        except Exception as e:
            self.error(f"Failed to save JSON file: {e}")
            return False

        self.success(f"Token saved to: {env_file}")
        self.success(f"Token metadata saved to: {json_file}")

        # Display token info (redacted for security)
        try:
            from ..utils import redact_sensitive_value
        except ImportError:
            # Fallback for when running as standalone script
            import sys
            from pathlib import Path

            utils_path = Path(__file__).parent.parent / "utils.py"
            if utils_path.exists():
                sys.path.insert(0, str(utils_path.parent))
                from utils import redact_sensitive_value
            else:
                # Simple fallback redaction function
                def redact_sensitive_value(value: str, show_chars: int = 8) -> str:
                    if not value or len(value) <= show_chars:
                        return "*" * len(value) if value else ""
                    return value[:show_chars] + "*" * (len(value) - show_chars)

        redacted_token = redact_sensitive_value(access_token, 8)
        self.logger.info(f"Access Token: {redacted_token}")
        if expires_in:
            self.logger.info(f"Expires in: {expires_in} seconds")
            if expires_at:
                expiry_time = datetime.fromisoformat(expires_at.replace("Z", "+00:00"))
                self.logger.info(f"Expires at: {expiry_time.strftime('%Y-%m-%d %H:%M:%S UTC')}")

        return True

    def generate_token_for_agent(
        self,
        agent_name: str,
        client_id: str = None,
        client_secret: str = None,
        keycloak_url: str = None,
        realm: str = "mcp-gateway",
        oauth_tokens_dir: str = None,
    ) -> bool:
        """Generate token for a single agent"""
        if oauth_tokens_dir is None:
            oauth_tokens_dir = os.path.join(
                os.path.dirname(os.path.dirname(os.path.dirname(__file__))), ".oauth-tokens"
            )

        # Load config from JSON if parameters not provided
        config = None
        if not all([client_id, client_secret, keycloak_url]):
            config = self.load_agent_config(agent_name, oauth_tokens_dir)
            if not config:
                return False

        # Use provided parameters or fall back to config
        if not client_id:
            client_id = config.get("client_id")
        if not client_secret:
            client_secret = config.get("client_secret")
        if not keycloak_url:
            keycloak_url = (
                config.get("keycloak_url") or config.get("gateway_url", "").split("/realms/")[0]
            )

        # Also try to get realm from config
        if config and realm == "mcp-gateway":
            config_realm = config.get("keycloak_realm") or config.get("realm")
            if config_realm:
                realm = config_realm

        # Validate required parameters
        if not client_id:
            self.error("CLIENT_ID is required. Provide via --client-id or in config file.")
            return False
        if not client_secret:
            self.error("CLIENT_SECRET is required. Provide via --client-secret or in config file.")
            return False
        if not keycloak_url:
            self.error("KEYCLOAK_URL is required. Provide via --keycloak-url or in config file.")
            return False

        self.logger.info(f"Requesting access token for agent: {agent_name}")

        # Get token from Keycloak
        token_data = self.get_token_from_keycloak(client_id, client_secret, keycloak_url, realm)
        if not token_data:
            return False

        self.success("Access token generated successfully!")

        # Save token files
        return self.save_token_files(
            agent_name, token_data, client_id, client_secret, keycloak_url, realm, oauth_tokens_dir
        )

    def find_agent_configs(self, oauth_tokens_dir: str) -> list[str]:
        """Find all agent-{}.json files, excluding agent-{}-token.json files"""
        if not os.path.exists(oauth_tokens_dir):
            self.warning(f"OAuth tokens directory not found: {oauth_tokens_dir}")
            return []

        # Find all agent-*.json files
        pattern = os.path.join(oauth_tokens_dir, "agent-*.json")
        all_files = glob.glob(pattern)

        # Filter out token files (agent-*-token.json)
        agent_configs = []
        for file_path in all_files:
            filename = os.path.basename(file_path)
            if not filename.endswith("-token.json"):
                # Use the full filename without extension as agent name
                agent_name = filename[:-5]  # Remove '.json' (5 chars)
                agent_configs.append(agent_name)

        return sorted(agent_configs)

    def generate_tokens_for_all_agents(
        self, oauth_tokens_dir: str = None, keycloak_url: str = None, realm: str = "mcp-gateway"
    ) -> bool:
        """Generate tokens for all agents found in .oauth-tokens directory"""
        if oauth_tokens_dir is None:
            oauth_tokens_dir = os.path.join(
                os.path.dirname(os.path.dirname(os.path.dirname(__file__))), ".oauth-tokens"
            )

        self.log(f"Searching for agent configs in: {oauth_tokens_dir}")

        agent_configs = self.find_agent_configs(oauth_tokens_dir)

        if not agent_configs:
            self.warning("No agent configuration files found")
            return True

        self.success(
            f"Found {len(agent_configs)} agent configuration(s): {', '.join(agent_configs)}"
        )

        success_count = 0
        total_count = len(agent_configs)

        for agent_name in agent_configs:
            self.logger.info("=" * 60)
            self.logger.info(f"Processing agent: {agent_name}")
            self.logger.info("=" * 60)

            try:
                if self.generate_token_for_agent(
                    agent_name,
                    keycloak_url=keycloak_url,
                    realm=realm,
                    oauth_tokens_dir=oauth_tokens_dir,
                ):
                    success_count += 1
                else:
                    self.error(f"Failed to generate token for agent: {agent_name}")
            except Exception as e:
                self.error(f"Exception while processing agent {agent_name}: {e}")

        self.logger.info("=" * 60)
        self.logger.info(f"Token generation complete: {success_count}/{total_count} successful")
        self.logger.info("=" * 60)

        return success_count == total_count


def main():
    """Main function"""
    parser = argparse.ArgumentParser(
        description="Generate OAuth2 access tokens for MCP agents using Keycloak",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Generate tokens for all agents in .oauth-tokens directory
  python generate_tokens.py --all-agents

  # Generate token for specific agent
  python generate_tokens.py --agent-name my-agent

  # Generate token with custom parameters
  python generate_tokens.py --agent-name my-agent --client-id custom-client --keycloak-url http://localhost:8080

  # Generate tokens for all agents with custom Keycloak URL
  python generate_tokens.py --all-agents --keycloak-url http://localhost:8080
        """,
    )

    parser.add_argument("--agent-name", type=str, help="Specific agent name to generate token for")
    parser.add_argument(
        "--all-agents",
        action="store_true",
        help="Generate tokens for all agents found in .oauth-tokens directory",
    )
    parser.add_argument("--client-id", type=str, help="OAuth2 client ID (overrides config file)")
    parser.add_argument(
        "--client-secret", type=str, help="OAuth2 client secret (overrides config file)"
    )
    parser.add_argument(
        "--keycloak-url", type=str, help="Keycloak server URL (overrides config file)"
    )
    parser.add_argument(
        "--realm", type=str, default="mcp-gateway", help="Keycloak realm (default: mcp-gateway)"
    )
    parser.add_argument(
        "--oauth-dir", type=str, help="OAuth tokens directory (default: ../../.oauth-tokens)"
    )
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")

    args = parser.parse_args()

    # Validate arguments
    if not args.all_agents and not args.agent_name:
        parser.error("Must specify either --all-agents or --agent-name")

    if args.all_agents and args.agent_name:
        parser.error("Cannot specify both --all-agents and --agent-name")

    # Initialize token generator
    generator = TokenGenerator(verbose=args.verbose)

    # Determine oauth tokens directory
    oauth_tokens_dir = args.oauth_dir
    if oauth_tokens_dir is None:
        oauth_tokens_dir = os.path.join(
            os.path.dirname(os.path.dirname(os.path.dirname(__file__))), ".oauth-tokens"
        )

    try:
        if args.all_agents:
            # Generate tokens for all agents
            success = generator.generate_tokens_for_all_agents(
                oauth_tokens_dir=oauth_tokens_dir, keycloak_url=args.keycloak_url, realm=args.realm
            )
        else:
            # Generate token for specific agent
            success = generator.generate_token_for_agent(
                agent_name=args.agent_name,
                client_id=args.client_id,
                client_secret=args.client_secret,
                keycloak_url=args.keycloak_url,
                realm=args.realm,
                oauth_tokens_dir=oauth_tokens_dir,
            )

        sys.exit(0 if success else 1)

    except KeyboardInterrupt:
        generator.warning("Operation interrupted by user")
        sys.exit(1)
    except Exception as e:
        generator.error(f"Unexpected error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
