from datetime import datetime
from typing import Dict, Any
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
import logging

from ..models import (
    ActionBatch,
    Investigation,
    Integration,
    IntegrationProvider,
    ActionStatus,
)
from .salesforce_gateway import SalesforceGateway
from .audit_service import AuditService

log = logging.getLogger(__name__)


class RollbackService:
    """Service for rolling back field updates"""

    def __init__(self, db: AsyncSession):
        self.db = db

    async def rollback_action_batch(
        self, batch_id: str, user, account_id: str, justification: str = None
    ) -> Dict[str, Any]:
        """Rollback field updates from an action batch"""

        # Get the batch scoped to requesting account
        stmt_batch = (
            select(ActionBatch, Investigation.account_id)
            .join(Investigation, Investigation.id == ActionBatch.investigation_id)
            .where(ActionBatch.id == batch_id)
            .where(Investigation.account_id == str(account_id))
            .limit(1)
        )
        batch_row = (await self.db.execute(stmt_batch)).first()
        if not batch_row:
            raise ValueError(f"Action batch {batch_id} not found")
        batch = batch_row[0]
        batch_account_id = str(batch_row[1])

        if batch.status != ActionStatus.APPLIED:
            raise ValueError("Can only rollback applied batches")

        # Get Salesforce integration for the account
        stmt = select(Integration).where(
            Integration.account_id == batch_account_id,
            Integration.provider == IntegrationProvider.SALESFORCE,
        )
        result = await self.db.execute(stmt)
        integration = result.scalar_one_or_none()

        if not integration:
            raise ValueError("Salesforce integration not found")

        # Create Salesforce gateway
        sf_client = await SalesforceGateway.for_tenant(integration.account_id, self.db)

        # Define callback to save token updates
        async def save_integration(updated_integration):
            await self.db.merge(updated_integration)
            await self.db.commit()
            await self.db.refresh(updated_integration)

        rollback_results = {"rolled_back": [], "skipped": [], "errors": []}

        # Process each change snapshot from the preview (created during apply)
        changes = batch.preview.get("changes", [])

        for change in changes:
            result = await self._rollback_field(sf_client, change, save_integration)

            if result["status"] == "rolled_back":
                rollback_results["rolled_back"].append(result)

                # Log audit event
                audit_service = AuditService(self.db)
                await audit_service.log_action(
                    actor=user,
                    action="rollback_field",
                    object_type=change["object_type"],
                    record_ids=[change["record_id"]],
                    changes={
                        f"{change['object_type']}:{change['record_id']}": {
                            change["field"]: {
                                "before": change[
                                    "after"
                                ],  # Current value (what was set)
                                "after": change["before"],  # Restoring to original
                            }
                        }
                    },
                    action_batch_id=batch_id,
                    justification=f"Rollback of batch {batch_id}: {justification or 'No justification provided'}",
                )
            elif result["status"] == "skipped":
                rollback_results["skipped"].append(result)
            else:
                rollback_results["errors"].append(result)

        rolled_count = len(rollback_results["rolled_back"])
        skipped_count = len(rollback_results["skipped"])
        error_count = len(rollback_results["errors"])
        if error_count > 0 and rolled_count == 0:
            rollback_status = "rollback_failed"
        elif error_count > 0 or skipped_count > 0:
            rollback_status = "partially_rolled_back"
        else:
            rollback_status = "rolled_back"

        # Persist rollback status and counts on the batch record.
        batch.status = rollback_status
        preview = dict(batch.preview or {})
        preview["rollback_summary"] = {
            "status": rollback_status,
            "total_changes": len(changes),
            "rolled_back_count": rolled_count,
            "skipped_count": skipped_count,
            "error_count": error_count,
            "completed_at": datetime.utcnow().isoformat(),
        }
        batch.preview = preview
        await self.db.commit()

        rollback_results["status"] = rollback_status
        return rollback_results

    async def _rollback_field(self, sf_client, change: Dict, save_callback) -> Dict:
        """Rollback a single field with compare-and-set"""
        try:
            # Get current value from Salesforce
            current = await sf_client.get_record(
                change["object_type"],
                change["record_id"],
                [change["field"]],
            )

            current_value = current.get(change["field"])
            expected_value = change["after"]  # This is what we set it to
            original_value = change["before"]  # This is what we're restoring to

            # For safety, we only rollback Account__c field for now
            if change["field"] != "Account__c":
                return {
                    "status": "skipped",
                    "record_id": change["record_id"],
                    "field": change["field"],
                    "reason": "Only Account__c field rollbacks are supported",
                }

            if current_value == expected_value:
                # Safe to rollback - restore to original value
                await sf_client.update(
                    change["object_type"],
                    change["record_id"],
                    {change["field"]: original_value},
                )

                log.info(
                    f"Rolled back {change['object_type']} {change['record_id']}: {change['field']} = {original_value}"
                )

                return {
                    "status": "rolled_back",
                    "record_id": change["record_id"],
                    "field": change["field"],
                    "previous_value": expected_value,
                    "restored_value": original_value,
                }
            else:
                # Value has changed since we last updated it
                return {
                    "status": "skipped",
                    "record_id": change["record_id"],
                    "field": change["field"],
                    "reason": f"Field value has been modified. Expected: {expected_value}, Found: {current_value}",
                }

        except Exception as e:
            log.error(
                f"Rollback error for {change.get('object_type')} {change.get('record_id')}: {e}"
            )
            return {
                "status": "error",
                "record_id": change.get("record_id"),
                "field": change.get("field"),
                "error": str(e),
            }
