"""Scraper worker for BYO Apify jobs with lease-based processing."""

from __future__ import annotations

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

from sqlalchemy import select, update, func
from sqlalchemy.ext.asyncio import AsyncSession

from .. import settings
from ..models import ScraperJob, ScraperJobStatus
from ..services.apify_adapter import ApifyScraperService
from ..services.scraper_blob_storage import ScraperBlobStorage
from ..services.secret_storage import get_saas_secret_backend

logger = logging.getLogger(__name__)


@dataclass
class WorkerConfig:
    lease_duration_seconds: int = 120
    chunk_size: int = 25
    max_errors_before_fail: int = 100
    progress_log_interval: int = 10


class LeaseNotAcquiredError(Exception):
    pass


class JobNotFoundError(Exception):
    pass


class ScraperWorker:
    def __init__(
        self,
        db: AsyncSession,
        job_id: UUID,
        task_id: str,
        config: Optional[WorkerConfig] = None,
    ) -> None:
        self.db = db
        self.job_id = job_id
        self.task_id = task_id
        self.config = config or WorkerConfig(
            lease_duration_seconds=settings.SCRAPER_LEASE_SECONDS,
            chunk_size=settings.SCRAPER_CHUNK_SIZE,
        )
        self._job: Optional[ScraperJob] = None
        self.blob_storage = ScraperBlobStorage()

    @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) -> ScraperJob:
        if self._job is None:
            result = await self.db.execute(
                select(ScraperJob).where(ScraperJob.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:
        now = datetime.now(timezone.utc)
        lease_expires = now + timedelta(seconds=self.config.lease_duration_seconds)
        result = await self.db.execute(
            update(ScraperJob)
            .where(
                ScraperJob.id == self.job_id,
                (ScraperJob.lease_expires_at.is_(None)) |
                (ScraperJob.lease_expires_at < now),
            )
            .values(
                lease_expires_at=lease_expires,
                lease_holder=self.task_id,
                status=ScraperJobStatus.RUNNING.value,
                started_at=func.coalesce(ScraperJob.started_at, now),
            )
            .returning(ScraperJob.id)
        )
        acquired = result.rowcount > 0
        if acquired:
            await self.db.commit()
            self._job = None
            logger.info("[SCRAPER_WORKER] Acquired lease for job %s", self.job_id)
        return acquired

    async def renew_lease(self) -> bool:
        now = datetime.now(timezone.utc)
        lease_expires = now + timedelta(seconds=self.config.lease_duration_seconds)
        result = await self.db.execute(
            update(ScraperJob)
            .where(
                ScraperJob.id == self.job_id,
                ScraperJob.lease_holder == self.task_id,
            )
            .values(lease_expires_at=lease_expires)
        )
        renewed = result.rowcount > 0
        if renewed:
            await self.db.commit()
        return renewed

    async def release_lease(self) -> None:
        await self.db.execute(
            update(ScraperJob)
            .where(
                ScraperJob.id == self.job_id,
                ScraperJob.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_scraped: int,
        rows_skipped: int,
        rows_errored: int,
    ) -> None:
        now = datetime.now(timezone.utc)
        lease_expires = now + timedelta(seconds=self.config.lease_duration_seconds)
        await self.db.execute(
            update(ScraperJob)
            .where(
                ScraperJob.id == self.job_id,
                ScraperJob.lease_holder == self.task_id,
            )
            .values(
                cursor=cursor,
                rows_processed=rows_processed,
                rows_scraped=rows_scraped,
                rows_skipped=rows_skipped,
                rows_errored=rows_errored,
                lease_expires_at=lease_expires,
            )
        )
        await self.db.commit()

    async def mark_completed(self) -> None:
        now = datetime.now(timezone.utc)
        await self.db.execute(
            update(ScraperJob)
            .where(ScraperJob.id == self.job_id)
            .values(
                status=ScraperJobStatus.COMPLETED.value,
                completed_at=now,
                lease_expires_at=None,
                lease_holder=None,
            )
        )
        await self.db.commit()
        logger.info("[SCRAPER_WORKER] Job %s completed", self.job_id)

    async def mark_failed(self, error_message: str) -> None:
        now = datetime.now(timezone.utc)
        await self.db.execute(
            update(ScraperJob)
            .where(ScraperJob.id == self.job_id)
            .values(
                status=ScraperJobStatus.FAILED.value,
                completed_at=now,
                lease_expires_at=None,
                lease_holder=None,
                error_message=error_message,
            )
        )
        await self.db.commit()
        logger.error("[SCRAPER_WORKER] Job %s failed: %s", self.job_id, error_message)

    async def process_chunk(self, rows: List[Dict[str, Any]]) -> Dict[str, Any]:
        if not await self.try_acquire_lease():
            raise LeaseNotAcquiredError(f"Could not acquire lease for job {self.job_id}")

        job = await self.get_job()
        if job.status == ScraperJobStatus.CANCELLED.value:
            await self.release_lease()
            return {"success": False, "cursor": job.cursor, "is_complete": True}

        backend = get_saas_secret_backend()
        secret = await backend.get_secret(
            db=self.db,
            integration="apify",
            key="api_token",
            workspace_id=job.workspace_id,
        )
        if not secret:
            await self.mark_failed("Apify not connected")
            raise RuntimeError("Apify not connected")

        service = ApifyScraperService(api_token=secret.value)
        actor = ApifyScraperService.get_actor(job.actor_type)
        required_keys = ApifyScraperService.required_keys(actor)

        rows_processed = job.rows_processed
        rows_scraped = job.rows_scraped
        rows_skipped = job.rows_skipped
        rows_errored = job.rows_errored
        cursor = job.cursor

        results: List[Dict[str, Any]] = []

        try:
            for idx, row in enumerate(rows):
                row_index = self._get_row_index(row, cursor + idx)
                inputs = service.extract_inputs(row, job.input_mapping or {}, actor)
                if not all(inputs.get(key) for key in required_keys):
                    rows_processed += 1
                    rows_skipped += 1
                    results.append(
                        {
                            "row_index": row_index,
                            "status": "skipped",
                            "error_message": "Missing required input",
                        }
                    )
                    continue
                try:
                    result = service.run_actor(inputs, actor, job.output_fields or [])
                    result["row_index"] = row_index
                    result["status"] = "scraped"
                    results.append(result)
                    rows_processed += 1
                    rows_scraped += 1
                except Exception as exc:
                    rows_processed += 1
                    rows_errored += 1
                    results.append(
                        {
                            "row_index": row_index,
                            "status": "error",
                            "error_message": str(exc),
                        }
                    )

                if rows_processed % self.config.progress_log_interval == 0:
                    logger.info(
                        "[SCRAPER_WORKER] Job %s progress: %s/%s rows",
                        self.job_id,
                        rows_processed,
                        job.rows_total,
                    )

            new_cursor = cursor + len(rows)

            # Persist results incrementally
            existing_results = await self.blob_storage.get_results(job.workspace_id, job.id)
            existing_results.extend(results)
            await self.blob_storage.store_results(job.workspace_id, job.id, existing_results)

            await self.update_progress(
                cursor=new_cursor,
                rows_processed=rows_processed,
                rows_scraped=rows_scraped,
                rows_skipped=rows_skipped,
                rows_errored=rows_errored,
            )

            is_complete = new_cursor >= job.rows_total
            if rows_errored > self.config.max_errors_before_fail:
                await self.mark_failed("Too many errors during scraping")
                return {
                    "success": False,
                    "cursor": new_cursor,
                    "rows_in_chunk": len(rows),
                    "is_complete": True,
                    "rows_processed": rows_processed,
                    "rows_scraped": rows_scraped,
                    "rows_skipped": rows_skipped,
                    "rows_errored": rows_errored,
                    "results": results,
                }

            if is_complete:
                await self.mark_completed()
            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_scraped": rows_scraped,
                "rows_skipped": rows_skipped,
                "rows_errored": rows_errored,
                "results": results,
            }

        except Exception as exc:
            logger.exception("[SCRAPER_WORKER] Job %s error: %s", self.job_id, exc)
            if rows_errored > self.config.max_errors_before_fail:
                await self.mark_failed(f"Too many errors: {exc}")
            else:
                await self.release_lease()
            raise


async def process_scraper_job(
    db: AsyncSession,
    job_id: UUID,
    task_id: str,
    rows: List[Dict[str, Any]],
) -> Dict[str, Any]:
    worker = ScraperWorker(db, job_id, task_id)
    return await worker.process_chunk(rows)
