"""
Salesforce Integration Core Module for FoundryMatch
==================================================
Comprehensive Salesforce integration with OAuth, API management, and data operations.
"""

import asyncio
import base64
import hashlib
import logging
import secrets
import time
import webbrowser
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Tuple
from urllib.parse import urlencode, parse_qs, urlparse
import threading
import queue
from http.server import HTTPServer, BaseHTTPRequestHandler
import socket

import pandas as pd
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

log = logging.getLogger(__name__)

# Constants
SALESFORCE_API_VERSION = "59.0"  # Winter '24
DEFAULT_TIMEOUT = 30
MAX_RETRIES = 3
CHUNK_SIZE = 10000  # For bulk operations
API_LIMIT_WARNING_THRESHOLD = 0.8  # 80% of daily limit


@dataclass
class SalesforceCredentials:
    """Salesforce OAuth credentials and connection info."""

    client_id: str
    client_secret: str
    redirect_uri: str = "http://localhost:8080/oauth/callback"
    instance_url: Optional[str] = None
    access_token: Optional[str] = None
    refresh_token: Optional[str] = None
    token_expires_at: Optional[datetime] = None
    org_id: Optional[str] = None
    user_id: Optional[str] = None
    username: Optional[str] = None
    is_sandbox: bool = False

    def is_expired(self) -> bool:
        """Check if access token is expired."""
        if not self.token_expires_at:
            return True
        return datetime.now() >= self.token_expires_at - timedelta(minutes=5)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "client_id": self.client_id,
            "client_secret": self.client_secret,
            "redirect_uri": self.redirect_uri,
            "instance_url": self.instance_url,
            "access_token": self.access_token,
            "refresh_token": self.refresh_token,
            "token_expires_at": self.token_expires_at.isoformat()
            if self.token_expires_at
            else None,
            "org_id": self.org_id,
            "user_id": self.user_id,
            "username": self.username,
            "is_sandbox": self.is_sandbox,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "SalesforceCredentials":
        """Create from dictionary."""
        creds = cls(
            client_id=data["client_id"],
            client_secret=data["client_secret"],
            redirect_uri=data.get(
                "redirect_uri", "http://localhost:8080/oauth/callback"
            ),
            instance_url=data.get("instance_url"),
            access_token=data.get("access_token"),
            refresh_token=data.get("refresh_token"),
            org_id=data.get("org_id"),
            user_id=data.get("user_id"),
            username=data.get("username"),
            is_sandbox=data.get("is_sandbox", False),
        )

        if data.get("token_expires_at"):
            creds.token_expires_at = datetime.fromisoformat(data["token_expires_at"])

        return creds


@dataclass
class SalesforceObject:
    """Represents a Salesforce object with metadata."""

    name: str
    label: str
    custom: bool = False
    fields: List["SalesforceField"] = field(default_factory=list)
    record_count: Optional[int] = None
    queryable: bool = True
    createable: bool = True
    updateable: bool = True
    deletable: bool = True


@dataclass
class SalesforceField:
    """Represents a Salesforce field with metadata."""

    name: str
    label: str
    type: str
    length: Optional[int] = None
    custom: bool = False
    nullable: bool = True
    unique: bool = False
    reference_to: Optional[List[str]] = None
    pick_list_values: Optional[List[str]] = None


@dataclass
class ApiLimits:
    """Track Salesforce API limits."""

    daily_api_requests_used: int = 0
    daily_api_requests_max: int = 0
    daily_bulk_api_requests_used: int = 0
    daily_bulk_api_requests_max: int = 0
    hourly_sync_report_runs_used: int = 0
    hourly_sync_report_runs_max: int = 0
    last_updated: datetime = field(default_factory=datetime.now)

    @property
    def api_usage_percentage(self) -> float:
        """Calculate API usage percentage."""
        if self.daily_api_requests_max == 0:
            return 0.0
        return (self.daily_api_requests_used / self.daily_api_requests_max) * 100

    @property
    def bulk_api_usage_percentage(self) -> float:
        """Calculate Bulk API usage percentage."""
        if self.daily_bulk_api_requests_max == 0:
            return 0.0
        return (
            self.daily_bulk_api_requests_used / self.daily_bulk_api_requests_max
        ) * 100

    def is_near_limit(self, threshold: float = API_LIMIT_WARNING_THRESHOLD) -> bool:
        """Check if near API limits."""
        return (
            self.api_usage_percentage >= threshold * 100
            or self.bulk_api_usage_percentage >= threshold * 100
        )


class OAuthCallbackHandler(BaseHTTPRequestHandler):
    """Handle OAuth callback from Salesforce."""

    def do_GET(self):
        """Handle GET request for OAuth callback."""
        try:
            parsed_url = urlparse(self.path)
            query_params = parse_qs(parsed_url.query)

            if "code" in query_params:
                # Success - authorization code received
                auth_code = query_params["code"][0]
                state = query_params.get("state", [None])[0]

                # Store result in server instance
                self.server.auth_result = {
                    "success": True,
                    "code": auth_code,
                    "state": state,
                }

                # Send success response
                self.send_response(200)
                self.send_header("Content-type", "text/html")
                self.end_headers()
                self.wfile.write(b"""
                <html>
                <head><title>Authorization Successful</title></head>
                <body>
                    <h1>Authorization Successful!</h1>
                    <p>You can now close this window and return to FoundryMatch.</p>
                    <script>setTimeout(function(){ window.close(); }, 3000);</script>
                </body>
                </html>
                """)

            elif "error" in query_params:
                # Error during authorization
                error = query_params["error"][0]
                error_description = query_params.get(
                    "error_description", ["Unknown error"]
                )[0]

                self.server.auth_result = {
                    "success": False,
                    "error": error,
                    "error_description": error_description,
                }

                # Send error response
                self.send_response(400)
                self.send_header("Content-type", "text/html")
                self.end_headers()
                self.wfile.write(
                    f"""
                <html>
                <head><title>Authorization Failed</title></head>
                <body>
                    <h1>Authorization Failed</h1>
                    <p>Error: {error}</p>
                    <p>Description: {error_description}</p>
                    <p>Please close this window and try again in FoundryMatch.</p>
                </body>
                </html>
                """.encode()
                )

        except Exception as e:
            log.error(f"Error handling OAuth callback: {e}", exc_info=True)
            self.send_response(500)
            self.send_header("Content-type", "text/html")
            self.end_headers()
            self.wfile.write(
                b"<html><body><h1>Internal Server Error</h1></body></html>"
            )

    def log_message(self, format, *args):
        """Suppress default logging."""
        pass


class SalesforceIntegration:
    """Main Salesforce integration class with comprehensive API support."""

    def __init__(self, credentials: Optional[SalesforceCredentials] = None):
        self.credentials = credentials
        self.session = self._create_session()
        self.api_limits = ApiLimits()
        self._oauth_server = None
        self._oauth_thread = None
        self._callback_queue = queue.Queue()

    def _create_session(self) -> requests.Session:
        """Create requests session with retry strategy."""
        session = requests.Session()

        # Configure retry strategy
        retry_strategy = Retry(
            total=MAX_RETRIES,
            status_forcelist=[429, 500, 502, 503, 504],
            allowed_methods=["HEAD", "GET", "OPTIONS"],
            backoff_factor=1,
        )

        adapter = HTTPAdapter(max_retries=retry_strategy)
        session.mount("http://", adapter)
        session.mount("https://", adapter)

        return session

    def _get_available_port(self, start_port: int = 8080) -> int:
        """Find an available port starting from start_port."""
        for port in range(start_port, start_port + 100):
            try:
                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                    s.bind(("localhost", port))
                    return port
            except OSError:
                continue
        raise RuntimeError("No available ports found")

    async def authenticate_oauth(
        self, client_id: str, client_secret: str, is_sandbox: bool = False
    ) -> bool:
        """
        Perform OAuth 2.0 authentication with Salesforce.

        Args:
            client_id: Salesforce Connected App Client ID
            client_secret: Salesforce Connected App Client Secret
            is_sandbox: Whether to connect to sandbox environment

        Returns:
            bool: True if authentication successful
        """
        try:
            # Find available port and update redirect URI
            port = self._get_available_port()
            redirect_uri = f"http://localhost:{port}/oauth/callback"

            # Initialize credentials
            self.credentials = SalesforceCredentials(
                client_id=client_id,
                client_secret=client_secret,
                redirect_uri=redirect_uri,
                is_sandbox=is_sandbox,
            )

            # Generate PKCE parameters
            code_verifier = (
                base64.urlsafe_b64encode(secrets.token_bytes(32))
                .decode("utf-8")
                .rstrip("=")
            )
            code_challenge = (
                base64.urlsafe_b64encode(
                    hashlib.sha256(code_verifier.encode("utf-8")).digest()
                )
                .decode("utf-8")
                .rstrip("=")
            )

            state = secrets.token_urlsafe(32)

            # Build authorization URL
            base_url = (
                "https://test.salesforce.com"
                if is_sandbox
                else "https://login.salesforce.com"
            )
            auth_params = {
                "response_type": "code",
                "client_id": client_id,
                "redirect_uri": redirect_uri,
                "scope": "full refresh_token offline_access",
                "state": state,
                "code_challenge": code_challenge,
                "code_challenge_method": "S256",
            }

            auth_url = f"{base_url}/services/oauth2/authorize?{urlencode(auth_params)}"

            # Start callback server
            callback_received = threading.Event()
            auth_result = {}

            def run_callback_server():
                try:
                    server = HTTPServer(("localhost", port), OAuthCallbackHandler)
                    server.auth_result = None
                    server.timeout = 1  # Short timeout for checking shutdown

                    log.info(f"OAuth callback server started on port {port}")

                    # Wait for callback or timeout
                    start_time = time.time()
                    while time.time() - start_time < 300:  # 5 minute timeout
                        server.handle_request()
                        if server.auth_result is not None:
                            auth_result.update(server.auth_result)
                            callback_received.set()
                            break

                    server.server_close()

                except Exception as e:
                    log.error(f"Error in OAuth callback server: {e}", exc_info=True)
                    auth_result.update(
                        {
                            "success": False,
                            "error": "server_error",
                            "error_description": str(e),
                        }
                    )
                    callback_received.set()

            # Start server in background thread
            server_thread = threading.Thread(target=run_callback_server, daemon=True)
            server_thread.start()

            # Open browser for authorization
            log.info("Opening browser for Salesforce authorization...")
            webbrowser.open(auth_url)

            # Wait for callback
            if callback_received.wait(timeout=300):  # 5 minute timeout
                if auth_result.get("success"):
                    # Exchange authorization code for tokens
                    auth_code = auth_result["code"]
                    return await self._exchange_auth_code(
                        auth_code, code_verifier, base_url
                    )
                else:
                    error = auth_result.get("error", "unknown_error")
                    error_desc = auth_result.get(
                        "error_description", "Unknown error occurred"
                    )
                    log.error(f"OAuth authorization failed: {error} - {error_desc}")
                    return False
            else:
                log.error("OAuth authorization timed out")
                return False

        except Exception as e:
            log.error(f"Error during OAuth authentication: {e}", exc_info=True)
            return False

    async def _exchange_auth_code(
        self, auth_code: str, code_verifier: str, base_url: str
    ) -> bool:
        """Exchange authorization code for access tokens."""
        try:
            token_url = f"{base_url}/services/oauth2/token"

            token_data = {
                "grant_type": "authorization_code",
                "client_id": self.credentials.client_id,
                "client_secret": self.credentials.client_secret,
                "redirect_uri": self.credentials.redirect_uri,
                "code": auth_code,
                "code_verifier": code_verifier,
            }

            # Use asyncio to run the blocking request in a thread
            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(
                None,
                lambda: self.session.post(
                    token_url, data=token_data, timeout=DEFAULT_TIMEOUT
                ),
            )

            if response.status_code == 200:
                token_info = response.json()

                # Update credentials
                self.credentials.access_token = token_info["access_token"]
                self.credentials.refresh_token = token_info.get("refresh_token")
                self.credentials.instance_url = token_info["instance_url"]

                # Calculate token expiration
                expires_in = token_info.get("expires_in", 3600)  # Default 1 hour
                self.credentials.token_expires_at = datetime.now() + timedelta(
                    seconds=expires_in
                )

                # Extract org and user info from ID URL
                id_url = token_info.get("id")
                if id_url:
                    # Get user info
                    user_info = await self._get_user_info(id_url)
                    if user_info:
                        self.credentials.org_id = user_info.get("organization_id")
                        self.credentials.user_id = user_info.get("user_id")
                        self.credentials.username = user_info.get("username")

                log.info(
                    f"Successfully authenticated with Salesforce org: {self.credentials.org_id}"
                )
                return True
            else:
                error_info = response.json() if response.content else {}
                log.error(
                    f"Token exchange failed: {response.status_code} - {error_info}"
                )
                return False

        except Exception as e:
            log.error(f"Error exchanging auth code: {e}", exc_info=True)
            return False

    async def _get_user_info(self, id_url: str) -> Optional[Dict[str, Any]]:
        """Get user information from Salesforce identity URL."""
        try:
            headers = {"Authorization": f"Bearer {self.credentials.access_token}"}

            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(
                None,
                lambda: self.session.get(
                    id_url, headers=headers, timeout=DEFAULT_TIMEOUT
                ),
            )

            if response.status_code == 200:
                return response.json()
            else:
                log.warning(f"Failed to get user info: {response.status_code}")
                return None

        except Exception as e:
            log.error(f"Error getting user info: {e}", exc_info=True)
            return None

    async def refresh_access_token(self) -> bool:
        """Refresh the access token using refresh token."""
        if not self.credentials or not self.credentials.refresh_token:
            log.error("No refresh token available")
            return False

        try:
            base_url = (
                "https://test.salesforce.com"
                if self.credentials.is_sandbox
                else "https://login.salesforce.com"
            )
            token_url = f"{base_url}/services/oauth2/token"

            token_data = {
                "grant_type": "refresh_token",
                "client_id": self.credentials.client_id,
                "client_secret": self.credentials.client_secret,
                "refresh_token": self.credentials.refresh_token,
            }

            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(
                None,
                lambda: self.session.post(
                    token_url, data=token_data, timeout=DEFAULT_TIMEOUT
                ),
            )

            if response.status_code == 200:
                token_info = response.json()

                # Update credentials
                self.credentials.access_token = token_info["access_token"]
                self.credentials.instance_url = token_info["instance_url"]

                # Calculate token expiration
                expires_in = token_info.get("expires_in", 3600)
                self.credentials.token_expires_at = datetime.now() + timedelta(
                    seconds=expires_in
                )

                log.info("Successfully refreshed access token")
                return True
            else:
                log.error(f"Token refresh failed: {response.status_code}")
                return False

        except Exception as e:
            log.error(f"Error refreshing token: {e}", exc_info=True)
            return False

    async def test_connection(self) -> Tuple[bool, str]:
        """Test the Salesforce connection."""
        try:
            if not self.credentials or not self.credentials.access_token:
                return False, "No valid credentials"

            # Check if token needs refresh
            if self.credentials.is_expired():
                if not await self.refresh_access_token():
                    return False, "Failed to refresh expired token"

            # Make a simple API call to test connection
            limits = await self.get_api_limits()
            if limits:
                return True, f"Connected to org: {self.credentials.org_id}"
            else:
                return False, "Failed to retrieve API limits"

        except Exception as e:
            log.error(f"Connection test failed: {e}", exc_info=True)
            return False, f"Connection error: {str(e)}"

    async def get_api_limits(self) -> Optional[ApiLimits]:
        """Get current API usage limits."""
        try:
            if not await self._ensure_valid_token():
                return None

            url = f"{self.credentials.instance_url}/services/data/v{SALESFORCE_API_VERSION}/limits"
            headers = {"Authorization": f"Bearer {self.credentials.access_token}"}

            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(
                None,
                lambda: self.session.get(url, headers=headers, timeout=DEFAULT_TIMEOUT),
            )

            if response.status_code == 200:
                limits_data = response.json()

                # Update API limits
                self.api_limits = ApiLimits(
                    daily_api_requests_used=limits_data.get("DailyApiRequests", {}).get(
                        "Remaining", 0
                    ),
                    daily_api_requests_max=limits_data.get("DailyApiRequests", {}).get(
                        "Max", 0
                    ),
                    daily_bulk_api_requests_used=limits_data.get(
                        "DailyBulkApiRequests", {}
                    ).get("Remaining", 0),
                    daily_bulk_api_requests_max=limits_data.get(
                        "DailyBulkApiRequests", {}
                    ).get("Max", 0),
                    last_updated=datetime.now(),
                )

                # Fix the calculation - Remaining is what's left, not what's used
                if self.api_limits.daily_api_requests_max > 0:
                    self.api_limits.daily_api_requests_used = (
                        self.api_limits.daily_api_requests_max
                        - self.api_limits.daily_api_requests_used
                    )

                if self.api_limits.daily_bulk_api_requests_max > 0:
                    self.api_limits.daily_bulk_api_requests_used = (
                        self.api_limits.daily_bulk_api_requests_max
                        - self.api_limits.daily_bulk_api_requests_used
                    )

                return self.api_limits
            else:
                log.error(f"Failed to get API limits: {response.status_code}")
                return None

        except Exception as e:
            log.error(f"Error getting API limits: {e}", exc_info=True)
            return None

    async def _ensure_valid_token(self) -> bool:
        """Ensure we have a valid access token."""
        if not self.credentials or not self.credentials.access_token:
            return False

        if self.credentials.is_expired():
            return await self.refresh_access_token()

        return True

    async def _execute_api_request(
        self, method: str, url: str, **kwargs
    ) -> requests.Response:
        """
        Execute API request with proper error handling and token management.

        Args:
            method: HTTP method (GET, POST, PATCH, PUT, DELETE)
            url: Full URL for the request
            **kwargs: Additional arguments for requests

        Returns:
            requests.Response object
        """
        if not await self._ensure_valid_token():
            raise RuntimeError("Invalid Salesforce token")

        # Ensure Authorization header is set
        headers = kwargs.get("headers", {})
        if "Authorization" not in headers:
            headers["Authorization"] = f"Bearer {self.credentials.access_token}"
            kwargs["headers"] = headers

        # Set default timeout if not provided
        if "timeout" not in kwargs:
            kwargs["timeout"] = DEFAULT_TIMEOUT

        loop = asyncio.get_event_loop()

        # Execute request based on method
        method = method.upper()
        if method == "GET":
            response = await loop.run_in_executor(
                None, lambda: self.session.get(url, **kwargs)
            )
        elif method == "POST":
            response = await loop.run_in_executor(
                None, lambda: self.session.post(url, **kwargs)
            )
        elif method == "PATCH":
            response = await loop.run_in_executor(
                None, lambda: self.session.patch(url, **kwargs)
            )
        elif method == "PUT":
            response = await loop.run_in_executor(
                None, lambda: self.session.put(url, **kwargs)
            )
        elif method == "DELETE":
            response = await loop.run_in_executor(
                None, lambda: self.session.delete(url, **kwargs)
            )
        else:
            raise ValueError(f"Unsupported HTTP method: {method}")

        # Log API usage for monitoring
        if hasattr(self, "api_limits"):
            # This would be more sophisticated in production
            # Could update API usage counters here
            pass

        return response

    async def get_objects(self, custom_only: bool = False) -> List[SalesforceObject]:
        """Get list of available Salesforce objects."""
        try:
            if not await self._ensure_valid_token():
                return []

            url = f"{self.credentials.instance_url}/services/data/v{SALESFORCE_API_VERSION}/sobjects"
            headers = {"Authorization": f"Bearer {self.credentials.access_token}"}

            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(
                None,
                lambda: self.session.get(url, headers=headers, timeout=DEFAULT_TIMEOUT),
            )

            if response.status_code == 200:
                data = response.json()
                objects = []

                for obj_info in data.get("sobjects", []):
                    if custom_only and not obj_info.get("custom", False):
                        continue

                    sf_obj = SalesforceObject(
                        name=obj_info["name"],
                        label=obj_info["label"],
                        custom=obj_info.get("custom", False),
                        queryable=obj_info.get("queryable", False),
                        createable=obj_info.get("createable", False),
                        updateable=obj_info.get("updateable", False),
                        deletable=obj_info.get("deletable", False),
                    )
                    objects.append(sf_obj)

                return sorted(objects, key=lambda x: x.label)
            else:
                log.error(f"Failed to get objects: {response.status_code}")
                return []

        except Exception as e:
            log.error(f"Error getting objects: {e}", exc_info=True)
            return []

    async def get_object_fields(self, object_name: str) -> List[SalesforceField]:
        """Get fields for a specific Salesforce object."""
        try:
            if not await self._ensure_valid_token():
                return []

            url = f"{self.credentials.instance_url}/services/data/v{SALESFORCE_API_VERSION}/sobjects/{object_name}/describe"
            headers = {"Authorization": f"Bearer {self.credentials.access_token}"}

            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(
                None,
                lambda: self.session.get(url, headers=headers, timeout=DEFAULT_TIMEOUT),
            )

            if response.status_code == 200:
                data = response.json()
                fields = []

                for field_info in data.get("fields", []):
                    sf_field = SalesforceField(
                        name=field_info["name"],
                        label=field_info["label"],
                        type=field_info["type"],
                        length=field_info.get("length"),
                        custom=field_info.get("custom", False),
                        nullable=field_info.get("nillable", True),
                        unique=field_info.get("unique", False),
                        reference_to=field_info.get("referenceTo"),
                        pick_list_values=[
                            pv["value"] for pv in field_info.get("picklistValues", [])
                        ],
                    )
                    fields.append(sf_field)

                return sorted(fields, key=lambda x: x.label)
            else:
                log.error(
                    f"Failed to get fields for {object_name}: {response.status_code}"
                )
                return []

        except Exception as e:
            log.error(f"Error getting fields for {object_name}: {e}", exc_info=True)
            return []

    async def query_data(
        self, soql: str, use_bulk: bool = False
    ) -> Optional[pd.DataFrame]:
        """Execute SOQL query and return results as DataFrame."""
        try:
            if not await self._ensure_valid_token():
                return None

            if use_bulk:
                return await self._bulk_query(soql)
            else:
                return await self._standard_query(soql)

        except Exception as e:
            log.error(f"Error executing query: {e}", exc_info=True)
            return None

    async def _standard_query(self, soql: str) -> Optional[pd.DataFrame]:
        """Execute standard SOQL query."""
        try:
            url = f"{self.credentials.instance_url}/services/data/v{SALESFORCE_API_VERSION}/query"
            headers = {"Authorization": f"Bearer {self.credentials.access_token}"}
            params = {"q": soql}

            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(
                None,
                lambda: self.session.get(
                    url, headers=headers, params=params, timeout=DEFAULT_TIMEOUT
                ),
            )

            if response.status_code == 200:
                data = response.json()
                records = data.get("records", [])

                # Handle pagination
                while not data.get("done", True):
                    next_url = (
                        f"{self.credentials.instance_url}{data['nextRecordsUrl']}"
                    )

                    next_response = await loop.run_in_executor(
                        None,
                        lambda: self.session.get(
                            next_url, headers=headers, timeout=DEFAULT_TIMEOUT
                        ),
                    )

                    if next_response.status_code == 200:
                        next_data = next_response.json()
                        records.extend(next_data.get("records", []))
                        data = next_data
                    else:
                        log.warning(
                            f"Failed to get next page: {next_response.status_code}"
                        )
                        break

                # Convert to DataFrame
                if records:
                    # Remove Salesforce metadata
                    clean_records = []
                    for record in records:
                        clean_record = {
                            k: v for k, v in record.items() if k != "attributes"
                        }
                        clean_records.append(clean_record)

                    return pd.DataFrame(clean_records)
                else:
                    return pd.DataFrame()
            else:
                log.error(f"Query failed: {response.status_code} - {response.text}")
                return None

        except Exception as e:
            log.error(f"Error in standard query: {e}", exc_info=True)
            return None

    async def _bulk_query(self, soql: str) -> Optional[pd.DataFrame]:
        """Execute bulk SOQL query for large datasets."""
        # This would implement Salesforce Bulk API 2.0
        # For now, fallback to standard query
        log.info("Bulk query not yet implemented, using standard query")
        return await self._standard_query(soql)

    def is_connected(self) -> bool:
        """Check if we have a valid connection to Salesforce."""
        return (
            self.credentials is not None
            and self.credentials.access_token is not None
            and not self.credentials.is_expired()
        )


# Add this class after SalesforceIntegration
class OrgLimitWatcher:
    """Monitor Salesforce API limits in real-time during workflows."""

    def __init__(self, sf_integration: SalesforceIntegration):
        self.sf_integration = sf_integration
        self.monitoring = False
        self.monitor_task = None

    async def start_monitoring(self):
        """Start background monitoring of API limits."""
        self.monitoring = True
        while self.monitoring:
            try:
                limits = await self.sf_integration.get_api_limits()
                if limits and limits.api_usage_percentage > 95:
                    log.error("🚨 API limit critical - aborting operations")
                    raise RuntimeError(
                        "API limit exceeded - workflow aborted for safety"
                    )
                elif limits and limits.api_usage_percentage > 85:
                    log.warning(f"⚠️ API usage high: {limits.api_usage_percentage:.1f}%")

                await asyncio.sleep(60)  # Check every minute
            except Exception as e:
                log.error(f"Limit monitoring error: {e}")
                await asyncio.sleep(60)

    def stop_monitoring(self):
        """Stop background monitoring."""
        self.monitoring = False

    async def check_limits_before_operation(self, estimated_calls: int) -> bool:
        """Check if we have enough API calls for an operation."""
        try:
            limits = await self.sf_integration.get_api_limits()
            if limits:
                remaining_calls = (
                    limits.daily_api_requests_max - limits.daily_api_requests_used
                )
                if estimated_calls > remaining_calls * 0.9:  # 90% safety margin
                    log.error(
                        f"Insufficient API calls: need {estimated_calls}, have {remaining_calls}"
                    )
                    return False
            return True
        except Exception as e:
            log.error(f"Error checking API limits: {e}")
            return False  # Fail safe


# Example usage and testing
async def test_salesforce_integration():
    """Test the Salesforce integration."""
    sf = SalesforceIntegration()

    # Test OAuth (you would need to provide real credentials)
    # success = await sf.authenticate_oauth("your_client_id", "your_client_secret", is_sandbox=True)
    # if success:
    #     print("Authentication successful!")
    #
    #     # Test connection
    #     connected, message = await sf.test_connection()
    #     print(f"Connection test: {connected} - {message}")
    #
    #     # Get objects
    #     objects = await sf.get_objects()
    #     print(f"Found {len(objects)} objects")
    #
    #     # Get Lead fields
    #     lead_fields = await sf.get_object_fields("Lead")
    #     print(f"Lead has {len(lead_fields)} fields")


if __name__ == "__main__":
    asyncio.run(test_salesforce_integration())
