"""
AIPTX Key Detector
==================

Unified API key detection from text and scan findings.

Detects:
- Cloud provider credentials (AWS, GCP, Azure)
- SaaS platform tokens (GitHub, Slack, Stripe, etc.)
- Database connection strings
- Private keys and certificates
"""

from __future__ import annotations

import re
import json
import logging
from dataclasses import dataclass
from typing import Any, Optional, TYPE_CHECKING

from .base import APIKeyInfo, KeyType, BaseKeyValidator

if TYPE_CHECKING:
    from aipt_v2.models.findings import Finding

logger = logging.getLogger(__name__)


@dataclass
class KeyPattern:
    """Key detection pattern configuration."""
    key_type: KeyType
    pattern: re.Pattern
    name: str
    description: str
    requires_context: bool = False  # Needs additional context (e.g., AWS secret needs access key)
    confidence: float = 0.9


# Comprehensive key detection patterns
KEY_PATTERNS: list[KeyPattern] = [
    # ==================== AWS ====================
    KeyPattern(
        key_type=KeyType.AWS,
        pattern=re.compile(r'AKIA[0-9A-Z]{16}'),
        name="aws_access_key",
        description="AWS Access Key ID",
        confidence=0.95,
    ),
    KeyPattern(
        key_type=KeyType.AWS,
        pattern=re.compile(r'(?:aws_secret_access_key|secret_key|secret)["\s:=]+["\']?([A-Za-z0-9/+=]{40})["\']?', re.I),
        name="aws_secret_key",
        description="AWS Secret Access Key",
        requires_context=True,
        confidence=0.85,
    ),
    KeyPattern(
        key_type=KeyType.AWS_SESSION,
        pattern=re.compile(r'ASIA[0-9A-Z]{16}'),
        name="aws_session_token",
        description="AWS Session Token (temporary credentials)",
        confidence=0.95,
    ),

    # ==================== GCP ====================
    KeyPattern(
        key_type=KeyType.GCP,
        pattern=re.compile(r'"type"\s*:\s*"service_account"', re.DOTALL),
        name="gcp_service_account",
        description="GCP Service Account JSON",
        confidence=0.95,
    ),
    KeyPattern(
        key_type=KeyType.GCP_API_KEY,
        pattern=re.compile(r'AIza[0-9A-Za-z_-]{35}'),
        name="gcp_api_key",
        description="GCP API Key",
        confidence=0.90,
    ),

    # ==================== Azure ====================
    KeyPattern(
        key_type=KeyType.AZURE,
        pattern=re.compile(r'(?:client_secret|azure_secret)["\s:=]+["\']?([A-Za-z0-9~._-]{34,})["\']?', re.I),
        name="azure_client_secret",
        description="Azure Client Secret",
        requires_context=True,
        confidence=0.80,
    ),
    KeyPattern(
        key_type=KeyType.AZURE_SAS,
        pattern=re.compile(r'sv=[0-9]{4}-[0-9]{2}-[0-9]{2}&s[prte]=.*?sig=[A-Za-z0-9%+/=]+'),
        name="azure_sas_token",
        description="Azure SAS Token",
        confidence=0.90,
    ),

    # ==================== GitHub ====================
    KeyPattern(
        key_type=KeyType.GITHUB,
        pattern=re.compile(r'ghp_[A-Za-z0-9]{36}'),
        name="github_pat_classic",
        description="GitHub Personal Access Token (Classic)",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.GITHUB_FINE_GRAINED,
        pattern=re.compile(r'github_pat_[A-Za-z0-9_]{22,}'),
        name="github_pat_fine_grained",
        description="GitHub Personal Access Token (Fine-grained)",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.GITHUB,
        pattern=re.compile(r'gho_[A-Za-z0-9]{36}'),
        name="github_oauth_token",
        description="GitHub OAuth Access Token",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.GITHUB,
        pattern=re.compile(r'ghu_[A-Za-z0-9]{36}'),
        name="github_user_to_server",
        description="GitHub User-to-Server Token",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.GITHUB,
        pattern=re.compile(r'ghs_[A-Za-z0-9]{36}'),
        name="github_server_to_server",
        description="GitHub Server-to-Server Token",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.GITHUB,
        pattern=re.compile(r'ghr_[A-Za-z0-9]{36}'),
        name="github_refresh_token",
        description="GitHub Refresh Token",
        confidence=0.98,
    ),

    # ==================== GitLab ====================
    KeyPattern(
        key_type=KeyType.GITLAB,
        pattern=re.compile(r'glpat-[A-Za-z0-9_-]{20,}'),
        name="gitlab_pat",
        description="GitLab Personal Access Token",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.GITLAB,
        pattern=re.compile(r'gloas-[A-Za-z0-9_-]{20,}'),
        name="gitlab_oauth",
        description="GitLab OAuth Token",
        confidence=0.98,
    ),

    # ==================== Slack ====================
    KeyPattern(
        key_type=KeyType.SLACK,
        pattern=re.compile(r'xoxb-[0-9]+-[0-9]+-[A-Za-z0-9]+'),
        name="slack_bot_token",
        description="Slack Bot Token",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.SLACK,
        pattern=re.compile(r'xoxp-[0-9]+-[0-9]+-[0-9]+-[a-f0-9]+'),
        name="slack_user_token",
        description="Slack User Token",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.SLACK,
        pattern=re.compile(r'xoxa-[0-9]+-[A-Za-z0-9]+'),
        name="slack_app_token",
        description="Slack App Token",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.SLACK_WEBHOOK,
        pattern=re.compile(r'https://hooks\.slack\.com/services/T[A-Z0-9]+/B[A-Z0-9]+/[A-Za-z0-9]+'),
        name="slack_webhook",
        description="Slack Webhook URL",
        confidence=0.95,
    ),

    # ==================== Stripe ====================
    KeyPattern(
        key_type=KeyType.STRIPE,
        pattern=re.compile(r'sk_live_[A-Za-z0-9]{24,}'),
        name="stripe_live_secret",
        description="Stripe Live Secret Key",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.STRIPE,
        pattern=re.compile(r'sk_test_[A-Za-z0-9]{24,}'),
        name="stripe_test_secret",
        description="Stripe Test Secret Key",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.STRIPE_RESTRICTED,
        pattern=re.compile(r'rk_live_[A-Za-z0-9]{24,}'),
        name="stripe_restricted_live",
        description="Stripe Restricted Live Key",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.STRIPE_RESTRICTED,
        pattern=re.compile(r'rk_test_[A-Za-z0-9]{24,}'),
        name="stripe_restricted_test",
        description="Stripe Restricted Test Key",
        confidence=0.98,
    ),

    # ==================== Twilio ====================
    KeyPattern(
        key_type=KeyType.TWILIO,
        pattern=re.compile(r'AC[a-f0-9]{32}'),
        name="twilio_account_sid",
        description="Twilio Account SID",
        confidence=0.95,
    ),
    KeyPattern(
        key_type=KeyType.TWILIO,
        pattern=re.compile(r'SK[a-f0-9]{32}'),
        name="twilio_api_key",
        description="Twilio API Key",
        confidence=0.95,
    ),

    # ==================== SendGrid ====================
    KeyPattern(
        key_type=KeyType.SENDGRID,
        pattern=re.compile(r'SG\.[A-Za-z0-9_-]{22}\.[A-Za-z0-9_-]{43}'),
        name="sendgrid_api_key",
        description="SendGrid API Key",
        confidence=0.98,
    ),

    # ==================== Mailgun ====================
    KeyPattern(
        key_type=KeyType.MAILGUN,
        pattern=re.compile(r'key-[a-f0-9]{32}'),
        name="mailgun_api_key",
        description="Mailgun API Key",
        confidence=0.95,
    ),

    # ==================== Discord ====================
    KeyPattern(
        key_type=KeyType.DISCORD,
        pattern=re.compile(r'[MN][A-Za-z0-9]{23,}\.[A-Za-z0-9_-]{6}\.[A-Za-z0-9_-]{27,}'),
        name="discord_bot_token",
        description="Discord Bot Token",
        confidence=0.90,
    ),
    KeyPattern(
        key_type=KeyType.DISCORD,
        pattern=re.compile(r'https://discord(?:app)?\.com/api/webhooks/[0-9]+/[A-Za-z0-9_-]+'),
        name="discord_webhook",
        description="Discord Webhook URL",
        confidence=0.95,
    ),

    # ==================== OpenAI / Anthropic ====================
    KeyPattern(
        key_type=KeyType.OPENAI,
        pattern=re.compile(r'sk-[A-Za-z0-9]{20}T3BlbkFJ[A-Za-z0-9]{20}'),
        name="openai_api_key",
        description="OpenAI API Key",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.OPENAI,
        pattern=re.compile(r'sk-proj-[A-Za-z0-9_-]{80,}'),
        name="openai_project_key",
        description="OpenAI Project API Key",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.ANTHROPIC,
        pattern=re.compile(r'sk-ant-[A-Za-z0-9_-]{90,}'),
        name="anthropic_api_key",
        description="Anthropic API Key",
        confidence=0.98,
    ),

    # ==================== Datadog ====================
    KeyPattern(
        key_type=KeyType.DATADOG,
        pattern=re.compile(r'[a-f0-9]{32}', re.I),
        name="datadog_api_key",
        description="Datadog API Key",
        requires_context=True,  # Too generic, needs context
        confidence=0.60,
    ),

    # ==================== Sentry ====================
    KeyPattern(
        key_type=KeyType.SENTRY,
        pattern=re.compile(r'https://[a-f0-9]{32}@[a-z0-9]+\.ingest\.sentry\.io/[0-9]+'),
        name="sentry_dsn",
        description="Sentry DSN",
        confidence=0.95,
    ),

    # ==================== Firebase ====================
    KeyPattern(
        key_type=KeyType.FIREBASE,
        pattern=re.compile(r'AAAA[A-Za-z0-9_-]{7}:[A-Za-z0-9_-]{140}'),
        name="firebase_cloud_messaging",
        description="Firebase Cloud Messaging Key",
        confidence=0.95,
    ),

    # ==================== Heroku ====================
    KeyPattern(
        key_type=KeyType.HEROKU,
        pattern=re.compile(r'[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}'),
        name="heroku_api_key",
        description="Heroku API Key (UUID format)",
        requires_context=True,
        confidence=0.70,
    ),

    # ==================== DigitalOcean ====================
    KeyPattern(
        key_type=KeyType.DIGITALOCEAN,
        pattern=re.compile(r'dop_v1_[a-f0-9]{64}'),
        name="digitalocean_pat",
        description="DigitalOcean Personal Access Token",
        confidence=0.98,
    ),
    KeyPattern(
        key_type=KeyType.DIGITALOCEAN,
        pattern=re.compile(r'doo_v1_[a-f0-9]{64}'),
        name="digitalocean_oauth",
        description="DigitalOcean OAuth Token",
        confidence=0.98,
    ),

    # ==================== NPM ====================
    KeyPattern(
        key_type=KeyType.NPM,
        pattern=re.compile(r'npm_[A-Za-z0-9]{36}'),
        name="npm_access_token",
        description="NPM Access Token",
        confidence=0.98,
    ),

    # ==================== PyPI ====================
    KeyPattern(
        key_type=KeyType.PYPI,
        pattern=re.compile(r'pypi-[A-Za-z0-9_-]{100,}'),
        name="pypi_api_token",
        description="PyPI API Token",
        confidence=0.98,
    ),

    # ==================== Docker Hub ====================
    KeyPattern(
        key_type=KeyType.DOCKER_HUB,
        pattern=re.compile(r'dckr_pat_[A-Za-z0-9_-]{20,}'),
        name="docker_hub_pat",
        description="Docker Hub Personal Access Token",
        confidence=0.98,
    ),

    # ==================== Private Keys ====================
    KeyPattern(
        key_type=KeyType.PRIVATE_KEY,
        pattern=re.compile(r'-----BEGIN (?:RSA |EC |OPENSSH |DSA )?PRIVATE KEY-----'),
        name="private_key_header",
        description="Private Key (RSA/EC/SSH/DSA)",
        confidence=0.95,
    ),
    KeyPattern(
        key_type=KeyType.PRIVATE_KEY,
        pattern=re.compile(r'-----BEGIN PGP PRIVATE KEY BLOCK-----'),
        name="pgp_private_key",
        description="PGP Private Key",
        confidence=0.95,
    ),

    # ==================== JWT ====================
    KeyPattern(
        key_type=KeyType.JWT,
        pattern=re.compile(r'eyJ[A-Za-z0-9_-]+\.eyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+'),
        name="jwt_token",
        description="JWT Token",
        confidence=0.85,
    ),

    # ==================== Database URLs ====================
    KeyPattern(
        key_type=KeyType.DATABASE_URL,
        pattern=re.compile(r'(?:mysql|postgres(?:ql)?|mariadb)://[^:]+:[^@]+@[^/]+/\w+'),
        name="sql_connection_string",
        description="SQL Database Connection String",
        confidence=0.90,
    ),
    KeyPattern(
        key_type=KeyType.MONGODB,
        pattern=re.compile(r'mongodb(?:\+srv)?://[^:]+:[^@]+@[^/]+'),
        name="mongodb_connection_string",
        description="MongoDB Connection String",
        confidence=0.90,
    ),
    KeyPattern(
        key_type=KeyType.REDIS,
        pattern=re.compile(r'redis://[^:]+:[^@]+@[^:]+:\d+'),
        name="redis_connection_string",
        description="Redis Connection String",
        confidence=0.90,
    ),

    # ==================== Generic Patterns ====================
    KeyPattern(
        key_type=KeyType.GENERIC_API_KEY,
        pattern=re.compile(r'(?:api[_-]?key|apikey)["\s:=]+["\']?([A-Za-z0-9_-]{20,})["\']?', re.I),
        name="generic_api_key",
        description="Generic API Key",
        requires_context=True,
        confidence=0.70,
    ),
    KeyPattern(
        key_type=KeyType.GENERIC_TOKEN,
        pattern=re.compile(r'(?:bearer|token|auth)["\s:=]+["\']?([A-Za-z0-9._-]{20,})["\']?', re.I),
        name="generic_token",
        description="Generic Bearer/Auth Token",
        requires_context=True,
        confidence=0.65,
    ),
]


class KeyDetector:
    """
    Unified API key detection from text and findings.

    Scans text for API key patterns and extracts them into
    structured APIKeyInfo objects for validation.
    """

    def __init__(self, include_low_confidence: bool = False):
        """
        Initialize detector.

        Args:
            include_low_confidence: Include patterns with confidence < 0.8
        """
        self.include_low_confidence = include_low_confidence
        self.patterns = self._filter_patterns()
        self._validators: dict[KeyType, BaseKeyValidator] = {}

    def _filter_patterns(self) -> list[KeyPattern]:
        """Filter patterns based on confidence threshold."""
        threshold = 0.6 if self.include_low_confidence else 0.8
        return [p for p in KEY_PATTERNS if p.confidence >= threshold]

    def register_validator(self, validator: BaseKeyValidator) -> None:
        """
        Register a validator for a key type.

        Args:
            validator: Validator instance
        """
        self._validators[validator.get_key_type()] = validator

    def get_validator(self, key_type: KeyType) -> Optional[BaseKeyValidator]:
        """
        Get validator for a key type.

        Args:
            key_type: Type of key

        Returns:
            Validator instance or None
        """
        return self._validators.get(key_type)

    def detect_from_text(self, text: str, source_url: str = None) -> list[APIKeyInfo]:
        """
        Detect API keys from raw text.

        Args:
            text: Text to scan for keys
            source_url: Optional source URL for context

        Returns:
            List of detected APIKeyInfo objects
        """
        if not text:
            return []

        detected: list[APIKeyInfo] = []
        seen_values: set[str] = set()

        for pattern in self.patterns:
            try:
                matches = pattern.pattern.finditer(text)
                for match in matches:
                    # Get the matched value (group 1 if capturing, else full match)
                    if match.groups():
                        value = match.group(1)
                    else:
                        value = match.group(0)

                    # Skip duplicates
                    if value in seen_values:
                        continue
                    seen_values.add(value)

                    # Skip very short matches
                    if len(value) < 10:
                        continue

                    key_info = APIKeyInfo(
                        key_type=pattern.key_type,
                        key_value=value,
                        detected_pattern=pattern.name,
                        confidence=pattern.confidence,
                        source_url=source_url,
                        metadata={
                            "pattern_name": pattern.name,
                            "pattern_description": pattern.description,
                            "requires_context": pattern.requires_context,
                        },
                    )

                    detected.append(key_info)
                    logger.debug(f"Detected {pattern.name}: {key_info.masked_value}")

            except Exception as e:
                logger.debug(f"Pattern {pattern.name} error: {e}")
                continue

        # Try to pair AWS access keys with secret keys
        detected = self._pair_aws_credentials(detected, text)

        return detected

    def detect_from_finding(self, finding: Any) -> list[APIKeyInfo]:
        """
        Detect API keys from a Finding object.

        Args:
            finding: Finding object with evidence/value

        Returns:
            List of detected APIKeyInfo objects
        """
        text_sources = []

        # Collect all text from finding
        if hasattr(finding, "evidence") and finding.evidence:
            text_sources.append(finding.evidence)
        if hasattr(finding, "value") and finding.value:
            text_sources.append(str(finding.value))
        if hasattr(finding, "description") and finding.description:
            text_sources.append(finding.description)
        if hasattr(finding, "request") and finding.request:
            text_sources.append(finding.request)
        if hasattr(finding, "response") and finding.response:
            text_sources.append(finding.response)

        combined_text = "\n".join(text_sources)
        source_url = getattr(finding, "url", None)

        keys = self.detect_from_text(combined_text, source_url)

        # Link back to finding
        for key in keys:
            key.source_finding = finding

        return keys

    def detect_from_json(self, json_data: dict | list | str) -> list[APIKeyInfo]:
        """
        Detect API keys from JSON data (e.g., scan results).

        Args:
            json_data: JSON data (dict, list, or JSON string)

        Returns:
            List of detected APIKeyInfo objects
        """
        if isinstance(json_data, str):
            try:
                json_data = json.loads(json_data)
            except json.JSONDecodeError:
                return self.detect_from_text(json_data)

        # Convert to string and scan
        json_str = json.dumps(json_data, indent=2)
        return self.detect_from_text(json_str)

    def detect_from_file(self, file_path: str) -> list[APIKeyInfo]:
        """
        Detect API keys from a file.

        Args:
            file_path: Path to file

        Returns:
            List of detected APIKeyInfo objects
        """
        try:
            with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
                content = f.read()

            keys = self.detect_from_text(content)

            # Add file source info
            for key in keys:
                key.source_file = file_path

            return keys

        except Exception as e:
            logger.error(f"Error reading file {file_path}: {e}")
            return []

    def _pair_aws_credentials(
        self, keys: list[APIKeyInfo], text: str
    ) -> list[APIKeyInfo]:
        """
        Try to pair AWS access keys with their corresponding secret keys.

        Args:
            keys: Detected keys
            text: Original text for context

        Returns:
            Updated list with paired credentials
        """
        aws_access_keys = [k for k in keys if k.key_type == KeyType.AWS and k.key_value.startswith("AKIA")]
        aws_secret_pattern = re.compile(r'[A-Za-z0-9/+=]{40}')

        for access_key in aws_access_keys:
            # Look for secret key near access key in text
            access_pos = text.find(access_key.key_value)
            if access_pos == -1:
                continue

            # Search in context around access key
            context_start = max(0, access_pos - 200)
            context_end = min(len(text), access_pos + 200)
            context = text[context_start:context_end]

            # Find potential secret keys
            for match in aws_secret_pattern.finditer(context):
                potential_secret = match.group(0)

                # Skip if it's the access key itself or too short
                if potential_secret == access_key.key_value:
                    continue

                # Verify it's not part of other patterns
                if any(p in potential_secret for p in ["AKIA", "ghp_", "sk_live"]):
                    continue

                # Add secret to access key
                access_key.secret_value = potential_secret
                logger.debug(f"Paired AWS secret with access key: {access_key.masked_value}")
                break

        return keys

    def get_supported_key_types(self) -> list[KeyType]:
        """Get list of all supported key types."""
        return list(set(p.key_type for p in self.patterns))

    def get_pattern_stats(self) -> dict:
        """Get statistics about loaded patterns."""
        return {
            "total_patterns": len(self.patterns),
            "key_types": len(self.get_supported_key_types()),
            "high_confidence": len([p for p in self.patterns if p.confidence >= 0.9]),
            "requires_context": len([p for p in self.patterns if p.requires_context]),
        }
