"""
Enrichment Worker for BYO Provider Locker.

Handles Cloud Tasks processing with:
- Lease-based exactly-once processing (compare-and-swap SQL)
- Cursor checkpointing for resume
- Chunk processing (configurable, default 50 rows)
- Job status updates

Usage:
    worker = EnrichmentWorker(db, job_id, task_id)
    await worker.process_chunk()
"""

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional, TYPE_CHECKING
from uuid import UUID

import logging
from sqlalchemy import select, update, func

logger = logging.getLogger(__name__)
from sqlalchemy.ext.asyncio import AsyncSession

from ..models import BYOEnrichmentJob, EnrichmentJobStatus, EnrichmentJobType
from .enrichment_adapters import (
    CascadingEnrichmentProvider,
    CompanyEnrichmentInput,
    ContactEnrichmentInput,
    EnrichmentStatus,
)
from .secret_cache import SecretCacheContext
from .. import settings

if TYPE_CHECKING:
    pass


@dataclass
class WorkerConfig:
    """Configuration for enrichment worker."""
    lease_duration_seconds: int = 120  # 2 minutes
    chunk_size: int = 50
    max_errors_before_fail: int = 100
    progress_log_interval: int = 10


class LeaseNotAcquiredError(Exception):
    """Raised when lease acquisition fails."""
    pass


class JobNotFoundError(Exception):
    """Raised when job doesn't exist."""
    pass


class EnrichmentWorker:
    """
    Enrichment worker that processes jobs with lease-based locking.

    Implements exactly-once processing via compare-and-swap SQL:
    1. Acquire lease with atomic UPDATE WHERE lease_expires_at IS NULL OR < NOW()
    2. Process chunk of rows
    3. Update cursor and renew lease
    4. Enqueue next task if more rows remain
    """

    def __init__(
        self,
        db: AsyncSession,
        job_id: UUID,
        task_id: str,
        config: Optional[WorkerConfig] = None,
    ):
        """
        Initialize enrichment worker.

        Args:
            db: Async database session
            job_id: ID of the enrichment job to process
            task_id: Unique identifier for this task (for lease ownership)
            config: Worker configuration
        """
        self.db = db
        self.job_id = job_id
        self.task_id = task_id
        self.config = config or WorkerConfig(
            lease_duration_seconds=settings.BYO_ENRICHMENT_LEASE_SECONDS,
            chunk_size=settings.BYO_ENRICHMENT_CHUNK_SIZE,
        )
        self._job: Optional[BYOEnrichmentJob] = None

    @staticmethod
    def _get_row_index(row: Dict[str, Any], fallback: int) -> int:
        value = row.get("__fm_row_index__")
        if value is None:
            return fallback
        try:
            return int(value)
        except (TypeError, ValueError):
            return fallback

    async def get_job(self) -> BYOEnrichmentJob:
        """Get the enrichment job, raising if not found."""
        if self._job is None:
            result = await self.db.execute(
                select(BYOEnrichmentJob).where(BYOEnrichmentJob.id == self.job_id)
            )
            self._job = result.scalar_one_or_none()
            if self._job is None:
                raise JobNotFoundError(f"Job {self.job_id} not found")
        return self._job

    async def try_acquire_lease(self) -> bool:
        """
        Try to acquire the job lease using compare-and-swap.

        Uses atomic SQL UPDATE with WHERE clause to ensure exactly-once.

        Returns:
            True if lease acquired, False if another worker owns it
        """
        now = datetime.now(timezone.utc)
        lease_expires = now + timedelta(seconds=self.config.lease_duration_seconds)

        # Atomic compare-and-swap: only acquire if no lease or expired
        result = await self.db.execute(
            update(BYOEnrichmentJob)
            .where(
                BYOEnrichmentJob.id == self.job_id,
                # Can acquire if: no lease, OR lease expired
                (BYOEnrichmentJob.lease_expires_at.is_(None)) |
                (BYOEnrichmentJob.lease_expires_at < now),
            )
            .values(
                lease_expires_at=lease_expires,
                lease_holder=self.task_id,
                status=EnrichmentJobStatus.RUNNING.value,
                started_at=func.coalesce(BYOEnrichmentJob.started_at, now),  # Only set if NULL
            )
            .returning(BYOEnrichmentJob.id)
        )

        acquired = result.rowcount > 0
        if acquired:
            await self.db.commit()
            # Refresh job state
            self._job = None
            logger.info(f"[WORKER] Acquired lease for job {self.job_id}")
        else:
            logger.debug(f"[WORKER] Lease not available for job {self.job_id}")

        return acquired

    async def renew_lease(self) -> bool:
        """
        Renew the lease if we still own it.

        Returns:
            True if renewed, False if lease was stolen (shouldn't happen)
        """
        now = datetime.now(timezone.utc)
        lease_expires = now + timedelta(seconds=self.config.lease_duration_seconds)

        result = await self.db.execute(
            update(BYOEnrichmentJob)
            .where(
                BYOEnrichmentJob.id == self.job_id,
                BYOEnrichmentJob.lease_holder == self.task_id,
            )
            .values(lease_expires_at=lease_expires)
        )

        renewed = result.rowcount > 0
        if renewed:
            await self.db.commit()
        else:
            logger.warning(f"[WORKER] Failed to renew lease for job {self.job_id}")

        return renewed

    async def release_lease(self) -> None:
        """Release the lease (called on completion or error)."""
        await self.db.execute(
            update(BYOEnrichmentJob)
            .where(
                BYOEnrichmentJob.id == self.job_id,
                BYOEnrichmentJob.lease_holder == self.task_id,
            )
            .values(
                lease_expires_at=None,
                lease_holder=None,
            )
        )
        await self.db.commit()

    async def update_progress(
        self,
        cursor: int,
        rows_processed: int,
        rows_enriched: int,
        rows_skipped: int,
        rows_errored: int,
        credits_consumed: int,
    ) -> None:
        """
        Update job progress and cursor.

        This also renews the lease to prevent timeout during long processing.
        """
        now = datetime.now(timezone.utc)
        lease_expires = now + timedelta(seconds=self.config.lease_duration_seconds)

        await self.db.execute(
            update(BYOEnrichmentJob)
            .where(
                BYOEnrichmentJob.id == self.job_id,
                BYOEnrichmentJob.lease_holder == self.task_id,
            )
            .values(
                cursor=cursor,
                rows_processed=rows_processed,
                rows_enriched=rows_enriched,
                rows_skipped=rows_skipped,
                rows_errored=rows_errored,
                credits_consumed=credits_consumed,
                lease_expires_at=lease_expires,
            )
        )
        await self.db.commit()

    async def mark_completed(self, summary: Dict[str, Any]) -> None:
        """Mark job as completed with summary."""
        now = datetime.now(timezone.utc)

        await self.db.execute(
            update(BYOEnrichmentJob)
            .where(BYOEnrichmentJob.id == self.job_id)
            .values(
                status=EnrichmentJobStatus.COMPLETED.value,
                completed_at=now,
                lease_expires_at=None,
                lease_holder=None,
                summary=summary,
            )
        )
        await self.db.commit()
        logger.info(f"[WORKER] Job {self.job_id} completed")

    async def mark_failed(self, error_message: str) -> None:
        """Mark job as failed with error message."""
        now = datetime.now(timezone.utc)

        await self.db.execute(
            update(BYOEnrichmentJob)
            .where(BYOEnrichmentJob.id == self.job_id)
            .values(
                status=EnrichmentJobStatus.FAILED.value,
                completed_at=now,
                lease_expires_at=None,
                lease_holder=None,
                error_message=error_message,
            )
        )
        await self.db.commit()
        logger.error(f"[WORKER] Job {self.job_id} failed: {error_message}")

    async def process_chunk(
        self,
        rows: List[Dict[str, Any]],
    ) -> Dict[str, Any]:
        """
        Process a chunk of rows.

        This is the main entry point called by the Cloud Tasks handler.

        Args:
            rows: List of row data to enrich (already sliced by cursor)

        Returns:
            Dict with processing results and next cursor
        """
        # Try to acquire lease
        if not await self.try_acquire_lease():
            raise LeaseNotAcquiredError(
                f"Could not acquire lease for job {self.job_id} - "
                "another worker may be processing"
            )

        job = await self.get_job()
        providers = job.providers[:1] if job.providers else ["foundrygraph"]
        rows_processed = job.rows_processed
        rows_enriched = job.rows_enriched
        rows_skipped = job.rows_skipped
        rows_errored = job.rows_errored
        credits_consumed = job.credits_consumed
        cursor = job.cursor

        try:
            # Setup providers with secret cache
            provider_set = {p for p in providers if p in {"fullenrich", "pdl"}}
            async with SecretCacheContext(
                self.db, job.workspace_id, provider_set
            ) as cache:
                async with CascadingEnrichmentProvider(
                    providers=providers,
                    secret_cache=cache,
                    db=self.db,
                    account_id=job.workspace_id,
                ) as provider:
                    # Prepare inputs
                    input_mapping = job.input_mapping or {}
                    inputs = []
                    if job.enrichment_type == EnrichmentJobType.CONTACT.value:
                        first_col = input_mapping.get("firstname", "firstname")
                        last_col = input_mapping.get("lastname", "lastname")
                        domain_col = input_mapping.get("domain", "domain")
                        company_col = input_mapping.get("company_name", "company_name")
                        email_col = input_mapping.get("email", "email")
                        linkedin_col = input_mapping.get("linkedin_url", "linkedin_url")
                        for i, row in enumerate(rows):
                            inputs.append(ContactEnrichmentInput(
                                first_name=row.get(first_col),
                                last_name=row.get(last_col),
                                domain=row.get(domain_col),
                                company_name=row.get(company_col),
                                email=row.get(email_col),
                                linkedin_url=row.get(linkedin_col),
                                row_index=self._get_row_index(row, cursor + i),
                                custom=row,
                            ))
                        results = await provider.enrich_contact(inputs)
                    else:
                        domain_col = input_mapping.get("domain", "domain")
                        name_col = input_mapping.get("company_name", "company_name")
                        for i, row in enumerate(rows):
                            inputs.append(CompanyEnrichmentInput(
                                domain=row.get(domain_col),
                                company_name=row.get(name_col),
                                row_index=self._get_row_index(row, cursor + i),
                                custom=row,
                            ))
                        results = await provider.enrich_company(inputs)

                    # Count results
                    for result in results:
                        rows_processed += 1
                        if result.status == EnrichmentStatus.ENRICHED:
                            rows_enriched += 1
                            credits_consumed += 1
                        elif result.status == EnrichmentStatus.SKIPPED:
                            rows_skipped += 1
                            # No credit for skipped rows
                        elif result.status == EnrichmentStatus.NOT_FOUND:
                            # No credit for not found (provider didn't charge)
                            pass
                        elif result.status == EnrichmentStatus.ERROR:
                            rows_errored += 1
                            # No credit for errors (user not charged for failures)

                        # Log progress periodically
                        if rows_processed % self.config.progress_log_interval == 0:
                            logger.info(
                                f"[WORKER] Job {self.job_id} progress: "
                                f"{rows_processed}/{job.row_count} rows"
                            )

            # Update cursor
            new_cursor = cursor + len(rows)

            # Update progress
            await self.update_progress(
                cursor=new_cursor,
                rows_processed=rows_processed,
                rows_enriched=rows_enriched,
                rows_skipped=rows_skipped,
                rows_errored=rows_errored,
                credits_consumed=credits_consumed,
            )

            # Check if done
            is_complete = new_cursor >= job.row_count

            if is_complete:
                # Build summary
                summary = {
                    "rows_total": job.row_count,
                    "rows_enriched": rows_enriched,
                    "rows_skipped": rows_skipped,
                    "rows_errored": rows_errored,
                    "credits_consumed": credits_consumed,
                    "provider_breakdown": {},  # TODO: Track per-provider counts
                }
                await self.mark_completed(summary)
            else:
                await self.release_lease()

            return {
                "success": True,
                "cursor": new_cursor,
                "rows_in_chunk": len(rows),
                "is_complete": is_complete,
                "rows_processed": rows_processed,
                "rows_enriched": rows_enriched,
                "rows_skipped": rows_skipped,
                "rows_errored": rows_errored,
                "credits_consumed": credits_consumed,
                "results": [r.to_dict() for r in results],
            }

        except Exception as e:
            logger.error(f"[WORKER] Job {self.job_id} error: {e}", exc_info=True)
            # Check if we hit error limit
            if rows_errored > self.config.max_errors_before_fail:
                await self.mark_failed(f"Too many errors: {e}")
            else:
                await self.release_lease()
            raise


async def process_enrichment_job(
    db: AsyncSession,
    job_id: UUID,
    task_id: str,
    rows: List[Dict[str, Any]],
) -> Dict[str, Any]:
    """
    Process an enrichment job chunk.

    This is the entry point called by Cloud Tasks handler.

    Args:
        db: Async database session
        job_id: Job ID to process
        task_id: Unique task ID for lease ownership
        rows: Rows to process in this chunk

    Returns:
        Processing result with cursor and status
    """
    worker = EnrichmentWorker(db, job_id, task_id)
    return await worker.process_chunk(rows)


async def resume_stalled_jobs(db: AsyncSession) -> int:
    """
    Find and reset jobs with expired leases.

    This should be called periodically (e.g., every 5 minutes) to handle
    jobs where the worker crashed without releasing the lease.

    Args:
        db: Async database session

    Returns:
        Number of jobs reset
    """
    now = datetime.now(timezone.utc)

    # Find jobs with expired leases that are still running
    result = await db.execute(
        update(BYOEnrichmentJob)
        .where(
            BYOEnrichmentJob.status == EnrichmentJobStatus.RUNNING.value,
            BYOEnrichmentJob.lease_expires_at.isnot(None),
            BYOEnrichmentJob.lease_expires_at < now,
        )
        .values(
            status=EnrichmentJobStatus.QUEUED.value,  # Reset to queued for retry
            lease_expires_at=None,
            lease_holder=None,
        )
        .returning(BYOEnrichmentJob.id)
    )

    reset_count = result.rowcount
    if reset_count > 0:
        await db.commit()
        logger.warning(f"[WORKER] Reset {reset_count} stalled enrichment jobs")

    return reset_count
