"""
Salesforce Bulk API 2.0 Integration for FoundryMatch
====================================================
Efficient handling of large datasets using Salesforce Bulk API 2.0 with progress tracking.
"""

import asyncio
import csv
import io
import json
import logging
import tempfile
import time
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum
from pathlib import Path
from typing import Dict, List, Optional, Any, Tuple, Union, AsyncGenerator
import uuid

import pandas as pd
import requests

from .core import SalesforceIntegration

log = logging.getLogger(__name__)


class BulkJobState(Enum):
    """Bulk job states."""

    OPEN = "Open"
    IN_PROGRESS = "InProgress"
    JOB_COMPLETE = "JobComplete"
    FAILED = "Failed"
    ABORTED = "Aborted"


class BulkJobType(Enum):
    """Bulk job types."""

    QUERY = "query"
    INSERT = "insert"
    UPDATE = "update"
    UPSERT = "upsert"
    DELETE = "delete"


@dataclass
class BulkJobInfo:
    """Information about a bulk job."""

    job_id: str
    object_name: str
    operation: BulkJobType
    state: BulkJobState
    created_date: datetime
    records_processed: int = 0
    records_failed: int = 0
    total_processing_time: int = 0  # milliseconds
    api_version: str = "59.0"
    error_message: Optional[str] = None


@dataclass
class BulkQueryResult:
    """Result of a bulk query operation."""

    job_info: BulkJobInfo
    data: Optional[pd.DataFrame] = None
    total_records: int = 0
    success: bool = False
    error_message: Optional[str] = None
    execution_time: float = 0.0


@dataclass
class BulkOperationResult:
    """Result of a bulk operation (insert/update/upsert/delete)."""

    job_info: BulkJobInfo
    successful_records: int = 0
    failed_records: int = 0
    total_records: int = 0
    success_file: Optional[str] = None
    error_file: Optional[str] = None
    execution_time: float = 0.0


class SalesforceBulkAPI:
    """Salesforce Bulk API 2.0 client with comprehensive error handling."""

    def __init__(self, sf_integration: SalesforceIntegration):
        self.sf_integration = sf_integration
        self.api_version = "59.0"
        self.chunk_size = 10000  # Records per batch
        self.max_poll_attempts = 300  # 5 minutes with 1s intervals
        self.poll_interval = 1.0  # seconds

    async def bulk_query(self, soql: str, progress_callback=None) -> BulkQueryResult:
        """
        Execute a bulk query for large datasets.

        Args:
            soql: SOQL query string
            progress_callback: Optional callback for progress updates

        Returns:
            BulkQueryResult with data and job information
        """
        start_time = time.time()

        try:
            log.info(f"Starting bulk query: {soql[:100]}...")

            # Create bulk query job
            job_info = await self._create_bulk_job("query", soql=soql)

            if progress_callback:
                await progress_callback("Bulk query job created", 10)

            # Wait for job completion
            final_job_info = await self._wait_for_job_completion(
                job_info, progress_callback, start_progress=10, end_progress=80
            )

            if final_job_info.state != BulkJobState.JOB_COMPLETE:
                return BulkQueryResult(
                    job_info=final_job_info,
                    success=False,
                    error_message=f"Job failed with state: {final_job_info.state.value}",
                    execution_time=time.time() - start_time,
                )

            if progress_callback:
                await progress_callback("Downloading query results", 85)

            # Download results
            data = await self._download_query_results(final_job_info.job_id)

            if progress_callback:
                await progress_callback("Query completed successfully", 100)

            execution_time = time.time() - start_time

            return BulkQueryResult(
                job_info=final_job_info,
                data=data,
                total_records=len(data) if data is not None else 0,
                success=True,
                execution_time=execution_time,
            )

        except Exception as e:
            log.error(f"Bulk query failed: {e}", exc_info=True)
            execution_time = time.time() - start_time

            return BulkQueryResult(
                job_info=BulkJobInfo(
                    job_id="",
                    object_name="",
                    operation=BulkJobType.QUERY,
                    state=BulkJobState.FAILED,
                    created_date=datetime.now(),
                ),
                success=False,
                error_message=str(e),
                execution_time=execution_time,
            )

    async def bulk_upsert(
        self,
        object_name: str,
        data: pd.DataFrame,
        external_id_field: str,
        progress_callback=None,
    ) -> BulkOperationResult:
        """
        Perform bulk upsert operation.

        Args:
            object_name: Salesforce object name
            data: DataFrame with records to upsert
            external_id_field: Field to use for upsert matching
            progress_callback: Optional callback for progress updates

        Returns:
            BulkOperationResult with operation details
        """
        return await self._bulk_operation(
            BulkJobType.UPSERT,
            object_name,
            data,
            external_id_field=external_id_field,
            progress_callback=progress_callback,
        )

    async def bulk_update(
        self, object_name: str, data: pd.DataFrame, progress_callback=None
    ) -> BulkOperationResult:
        """
        Perform bulk update operation.

        Args:
            object_name: Salesforce object name
            data: DataFrame with records to update (must include Id field)
            progress_callback: Optional callback for progress updates

        Returns:
            BulkOperationResult with operation details
        """
        if "Id" not in data.columns:
            raise ValueError("Update operations require 'Id' field in data")

        return await self._bulk_operation(
            BulkJobType.UPDATE, object_name, data, progress_callback=progress_callback
        )

    async def bulk_insert(
        self, object_name: str, data: pd.DataFrame, progress_callback=None
    ) -> BulkOperationResult:
        """
        Perform bulk insert operation.

        Args:
            object_name: Salesforce object name
            data: DataFrame with records to insert
            progress_callback: Optional callback for progress updates

        Returns:
            BulkOperationResult with operation details
        """
        return await self._bulk_operation(
            BulkJobType.INSERT, object_name, data, progress_callback=progress_callback
        )

    async def bulk_delete(
        self, object_name: str, record_ids: List[str], progress_callback=None
    ) -> BulkOperationResult:
        """
        Perform bulk delete operation.

        Args:
            object_name: Salesforce object name
            record_ids: List of record IDs to delete
            progress_callback: Optional callback for progress updates

        Returns:
            BulkOperationResult with operation details
        """
        # Convert IDs to DataFrame
        data = pd.DataFrame({"Id": record_ids})

        return await self._bulk_operation(
            BulkJobType.DELETE, object_name, data, progress_callback=progress_callback
        )

    async def _bulk_operation(
        self,
        operation: BulkJobType,
        object_name: str,
        data: pd.DataFrame,
        external_id_field: Optional[str] = None,
        progress_callback=None,
    ) -> BulkOperationResult:
        """Execute a bulk operation with comprehensive error handling."""
        start_time = time.time()

        try:
            log.info(
                f"Starting bulk {operation.value} on {object_name} with {len(data)} records"
            )

            # Create bulk job
            job_info = await self._create_bulk_job(
                operation.value, object_name, external_id_field
            )

            if progress_callback:
                await progress_callback(f"Created bulk {operation.value} job", 5)

            # Upload data
            await self._upload_job_data(job_info.job_id, data)

            if progress_callback:
                await progress_callback("Data uploaded, processing...", 20)

            # Close job to start processing
            await self._close_bulk_job(job_info.job_id)

            if progress_callback:
                await progress_callback("Job started, waiting for completion", 25)

            # Wait for completion
            final_job_info = await self._wait_for_job_completion(
                job_info, progress_callback, start_progress=25, end_progress=85
            )

            if progress_callback:
                await progress_callback("Downloading results", 90)

            # Download results
            success_file, error_file = await self._download_operation_results(
                final_job_info.job_id
            )

            if progress_callback:
                await progress_callback("Operation completed", 100)

            execution_time = time.time() - start_time

            return BulkOperationResult(
                job_info=final_job_info,
                successful_records=final_job_info.records_processed
                - final_job_info.records_failed,
                failed_records=final_job_info.records_failed,
                total_records=len(data),
                success_file=success_file,
                error_file=error_file,
                execution_time=execution_time,
            )

        except Exception as e:
            log.error(f"Bulk {operation.value} failed: {e}", exc_info=True)
            execution_time = time.time() - start_time

            return BulkOperationResult(
                job_info=BulkJobInfo(
                    job_id="",
                    object_name=object_name,
                    operation=operation,
                    state=BulkJobState.FAILED,
                    created_date=datetime.now(),
                ),
                failed_records=len(data),
                total_records=len(data),
                execution_time=execution_time,
            )

    async def _create_bulk_job(
        self,
        operation: str,
        object_name: Optional[str] = None,
        external_id_field: Optional[str] = None,
        soql: Optional[str] = None,
    ) -> BulkJobInfo:
        """Create a new bulk job."""
        if not await self.sf_integration._ensure_valid_token():
            raise RuntimeError("Invalid Salesforce token")

        url = f"{self.sf_integration.credentials.instance_url}/services/data/v{self.api_version}/jobs/{operation}"

        headers = {
            "Authorization": f"Bearer {self.sf_integration.credentials.access_token}",
            "Content-Type": "application/json",
        }

        job_data = {"operation": operation, "contentType": "CSV", "lineEnding": "LF"}

        if object_name:
            job_data["object"] = object_name

        if external_id_field:
            job_data["externalIdFieldName"] = external_id_field

        if soql:
            job_data["query"] = soql

        loop = asyncio.get_event_loop()
        response = await loop.run_in_executor(
            None,
            lambda: self.sf_integration.session.post(
                url, headers=headers, json=job_data, timeout=30
            ),
        )

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

            return BulkJobInfo(
                job_id=job_response["id"],
                object_name=job_response.get("object", ""),
                operation=BulkJobType(operation),
                state=BulkJobState(job_response["state"]),
                created_date=datetime.fromisoformat(
                    job_response["createdDate"].replace("Z", "+00:00")
                ),
                api_version=self.api_version,
            )
        else:
            error_msg = (
                f"Failed to create bulk job: {response.status_code} - {response.text}"
            )
            log.error(error_msg)
            raise RuntimeError(error_msg)

    async def _upload_job_data(self, job_id: str, data: pd.DataFrame):
        """Upload CSV data to a bulk job."""
        if not await self.sf_integration._ensure_valid_token():
            raise RuntimeError("Invalid Salesforce token")

        url = f"{self.sf_integration.credentials.instance_url}/services/data/v{self.api_version}/jobs/ingest/{job_id}/batches"

        headers = {
            "Authorization": f"Bearer {self.sf_integration.credentials.access_token}",
            "Content-Type": "text/csv",
        }

        # Convert DataFrame to CSV string
        csv_buffer = io.StringIO()
        data.to_csv(csv_buffer, index=False, quoting=csv.QUOTE_ALL)
        csv_data = csv_buffer.getvalue()

        loop = asyncio.get_event_loop()
        response = await loop.run_in_executor(
            None,
            lambda: self.sf_integration.session.put(
                url, headers=headers, data=csv_data.encode("utf-8"), timeout=300
            ),
        )

        if response.status_code != 200:
            error_msg = (
                f"Failed to upload job data: {response.status_code} - {response.text}"
            )
            log.error(error_msg)
            raise RuntimeError(error_msg)

    async def _close_bulk_job(self, job_id: str):
        """Close a bulk job to start processing."""
        if not await self.sf_integration._ensure_valid_token():
            raise RuntimeError("Invalid Salesforce token")

        url = f"{self.sf_integration.credentials.instance_url}/services/data/v{self.api_version}/jobs/ingest/{job_id}"

        headers = {
            "Authorization": f"Bearer {self.sf_integration.credentials.access_token}",
            "Content-Type": "application/json",
        }

        close_data = {"state": "UploadComplete"}

        loop = asyncio.get_event_loop()
        response = await loop.run_in_executor(
            None,
            lambda: self.sf_integration.session.patch(
                url, headers=headers, json=close_data, timeout=30
            ),
        )

        if response.status_code != 200:
            error_msg = (
                f"Failed to close bulk job: {response.status_code} - {response.text}"
            )
            log.error(error_msg)
            raise RuntimeError(error_msg)

    async def _wait_for_job_completion(
        self,
        job_info: BulkJobInfo,
        progress_callback=None,
        start_progress: int = 0,
        end_progress: int = 100,
    ) -> BulkJobInfo:
        """Wait for bulk job to complete with progress updates."""
        if not await self.sf_integration._ensure_valid_token():
            raise RuntimeError("Invalid Salesforce token")

        # Determine URL based on job type
        if job_info.operation == BulkJobType.QUERY:
            url = f"{self.sf_integration.credentials.instance_url}/services/data/v{self.api_version}/jobs/query/{job_info.job_id}"
        else:
            url = f"{self.sf_integration.credentials.instance_url}/services/data/v{self.api_version}/jobs/ingest/{job_info.job_id}"

        headers = {
            "Authorization": f"Bearer {self.sf_integration.credentials.access_token}",
            "Content-Type": "application/json",
        }

        for attempt in range(self.max_poll_attempts):
            try:
                loop = asyncio.get_event_loop()
                response = await loop.run_in_executor(
                    None,
                    lambda: self.sf_integration.session.get(
                        url, headers=headers, timeout=30
                    ),
                )

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

                    # Update job info
                    job_info.state = BulkJobState(job_data["state"])
                    job_info.records_processed = job_data.get("recordsProcessed", 0)
                    job_info.records_failed = job_data.get("recordsFailed", 0)
                    job_info.total_processing_time = job_data.get(
                        "totalProcessingTime", 0
                    )

                    # Calculate progress within the specified range
                    if progress_callback:
                        progress_pct = start_progress + (
                            (attempt / self.max_poll_attempts)
                            * (end_progress - start_progress)
                        )
                        await progress_callback(
                            f"Job {job_info.state.value}: {job_info.records_processed} records processed",
                            progress_pct,
                        )

                    # Check if job is complete
                    if job_info.state in [
                        BulkJobState.JOB_COMPLETE,
                        BulkJobState.FAILED,
                        BulkJobState.ABORTED,
                    ]:
                        if job_info.state == BulkJobState.FAILED:
                            job_info.error_message = job_data.get(
                                "stateMessage", "Job failed"
                            )

                        log.info(
                            f"Bulk job {job_info.job_id} completed with state: {job_info.state.value}"
                        )
                        return job_info

                    # Wait before next poll
                    await asyncio.sleep(self.poll_interval)

                else:
                    log.warning(f"Failed to get job status: {response.status_code}")
                    await asyncio.sleep(self.poll_interval * 2)  # Longer wait on error

            except Exception as e:
                log.warning(f"Error polling job status: {e}")
                await asyncio.sleep(self.poll_interval * 2)

        # Timeout reached
        job_info.state = BulkJobState.FAILED
        job_info.error_message = "Job polling timeout"
        log.error(
            f"Bulk job {job_info.job_id} polling timeout after {self.max_poll_attempts} attempts"
        )

        return job_info

    async def _download_query_results(self, job_id: str) -> Optional[pd.DataFrame]:
        """Download results from a completed query job."""
        if not await self.sf_integration._ensure_valid_token():
            raise RuntimeError("Invalid Salesforce token")

        url = f"{self.sf_integration.credentials.instance_url}/services/data/v{self.api_version}/jobs/query/{job_id}/results"

        headers = {
            "Authorization": f"Bearer {self.sf_integration.credentials.access_token}",
            "Accept": "text/csv",
        }

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

        if response.status_code == 200:
            # Read CSV data directly into DataFrame
            csv_content = response.content.decode("utf-8")
            csv_buffer = io.StringIO(csv_content)

            try:
                df = pd.read_csv(csv_buffer)
                log.info(f"Downloaded {len(df)} records from query job {job_id}")
                return df
            except Exception as e:
                log.error(f"Error parsing CSV results: {e}")
                return None
        else:
            error_msg = f"Failed to download query results: {response.status_code} - {response.text}"
            log.error(error_msg)
            raise RuntimeError(error_msg)

    async def _download_operation_results(
        self, job_id: str
    ) -> Tuple[Optional[str], Optional[str]]:
        """Download success and error results from a completed operation job."""
        if not await self.sf_integration._ensure_valid_token():
            raise RuntimeError("Invalid Salesforce token")

        base_url = f"{self.sf_integration.credentials.instance_url}/services/data/v{self.api_version}/jobs/ingest/{job_id}"

        headers = {
            "Authorization": f"Bearer {self.sf_integration.credentials.access_token}",
            "Accept": "text/csv",
        }

        success_file = None
        error_file = None

        # Download successful records
        try:
            success_url = f"{base_url}/successfulResults"

            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(
                None,
                lambda: self.sf_integration.session.get(
                    success_url, headers=headers, timeout=300
                ),
            )

            if response.status_code == 200 and response.content:
                success_file = tempfile.NamedTemporaryFile(
                    mode="w", suffix="_success.csv", delete=False
                )
                success_file.write(response.content.decode("utf-8"))
                success_file.close()
                success_file = success_file.name
                log.info(f"Downloaded successful results to {success_file}")

        except Exception as e:
            log.warning(f"Error downloading successful results: {e}")

        # Download failed records
        try:
            error_url = f"{base_url}/failedResults"

            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(
                None,
                lambda: self.sf_integration.session.get(
                    error_url, headers=headers, timeout=300
                ),
            )

            if response.status_code == 200 and response.content:
                error_file = tempfile.NamedTemporaryFile(
                    mode="w", suffix="_errors.csv", delete=False
                )
                error_file.write(response.content.decode("utf-8"))
                error_file.close()
                error_file = error_file.name
                log.info(f"Downloaded error results to {error_file}")

        except Exception as e:
            log.warning(f"Error downloading error results: {e}")

        return success_file, error_file

    async def get_job_info(
        self, job_id: str, job_type: str = "ingest"
    ) -> Optional[BulkJobInfo]:
        """Get information about a bulk job."""
        if not await self.sf_integration._ensure_valid_token():
            return None

        url = f"{self.sf_integration.credentials.instance_url}/services/data/v{self.api_version}/jobs/{job_type}/{job_id}"

        headers = {
            "Authorization": f"Bearer {self.sf_integration.credentials.access_token}",
            "Content-Type": "application/json",
        }

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

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

                return BulkJobInfo(
                    job_id=job_data["id"],
                    object_name=job_data.get("object", ""),
                    operation=BulkJobType(job_data["operation"]),
                    state=BulkJobState(job_data["state"]),
                    created_date=datetime.fromisoformat(
                        job_data["createdDate"].replace("Z", "+00:00")
                    ),
                    records_processed=job_data.get("recordsProcessed", 0),
                    records_failed=job_data.get("recordsFailed", 0),
                    total_processing_time=job_data.get("totalProcessingTime", 0),
                    api_version=self.api_version,
                )
            else:
                log.error(f"Failed to get job info: {response.status_code}")
                return None

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

    async def abort_job(self, job_id: str, job_type: str = "ingest") -> bool:
        """Abort a running bulk job."""
        if not await self.sf_integration._ensure_valid_token():
            return False

        url = f"{self.sf_integration.credentials.instance_url}/services/data/v{self.api_version}/jobs/{job_type}/{job_id}"

        headers = {
            "Authorization": f"Bearer {self.sf_integration.credentials.access_token}",
            "Content-Type": "application/json",
        }

        abort_data = {"state": "Aborted"}

        try:
            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(
                None,
                lambda: self.sf_integration.session.patch(
                    url, headers=headers, json=abort_data, timeout=30
                ),
            )

            if response.status_code == 200:
                log.info(f"Successfully aborted job {job_id}")
                return True
            else:
                log.error(f"Failed to abort job: {response.status_code}")
                return False

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


# Helper functions for integration


async def execute_bulk_query_with_retry(
    bulk_api: SalesforceBulkAPI, soql: str, max_retries: int = 3, progress_callback=None
) -> BulkQueryResult:
    """Execute bulk query with retry logic."""

    for attempt in range(max_retries):
        try:
            if progress_callback:
                await progress_callback(f"Query attempt {attempt + 1}/{max_retries}", 0)

            result = await bulk_api.bulk_query(soql, progress_callback)

            if result.success:
                return result

            if attempt < max_retries - 1:
                log.warning(
                    f"Query attempt {attempt + 1} failed: {result.error_message}"
                )
                await asyncio.sleep(2**attempt)  # Exponential backoff

        except Exception as e:
            if attempt < max_retries - 1:
                log.warning(f"Query attempt {attempt + 1} failed with exception: {e}")
                await asyncio.sleep(2**attempt)
            else:
                raise

    # All attempts failed
    return BulkQueryResult(
        job_info=BulkJobInfo(
            job_id="",
            object_name="",
            operation=BulkJobType.QUERY,
            state=BulkJobState.FAILED,
            created_date=datetime.now(),
        ),
        success=False,
        error_message="All retry attempts failed",
    )


def estimate_bulk_api_benefits(record_count: int, field_count: int) -> Dict[str, Any]:
    """Estimate benefits of using Bulk API vs standard API."""

    # Standard API limits (approximate)
    standard_api_calls = record_count  # 1 call per record for most operations
    standard_time_estimate = record_count * 0.1  # 100ms per call (optimistic)

    # Bulk API estimates
    bulk_api_calls = 1  # Single job regardless of size
    bulk_time_estimate = max(60, record_count * 0.01)  # Minimum 1 minute overhead

    # Calculate savings
    api_call_savings = standard_api_calls - bulk_api_calls
    time_savings = standard_time_estimate - bulk_time_estimate

    # Recommendations
    recommended = record_count >= 200 or field_count >= 20

    return {
        "record_count": record_count,
        "field_count": field_count,
        "standard_api_calls": standard_api_calls,
        "bulk_api_calls": bulk_api_calls,
        "api_call_savings": api_call_savings,
        "estimated_time_savings_seconds": time_savings,
        "recommended": recommended,
        "reasoning": (
            f"Bulk API recommended for {record_count:,} records"
            if recommended
            else f"Standard API sufficient for {record_count:,} records"
        ),
    }


# Example usage and testing
async def test_bulk_operations():
    """Test bulk operations."""
    # This would use a real SalesforceIntegration instance
    # sf_integration = SalesforceIntegration()
    # bulk_api = SalesforceBulkAPI(sf_integration)

    # Test data
    # test_data = pd.DataFrame({
    #     'Name': [f'Test Account {i}' for i in range(1000)],
    #     'Type': ['Customer'] * 1000,
    #     'Industry': ['Technology'] * 1000
    # })

    # Bulk insert example
    # result = await bulk_api.bulk_insert('Account', test_data)
    # print(f"Bulk insert result: {result.successful_records} successful, {result.failed_records} failed")

    pass


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