"""Sequencer export worker."""

from __future__ import annotations

import asyncio
import json
import logging
import os
import re
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from g_gremlin.sequencer import (
    SequencerAuthError,
    SequencerError,
    SequencerRateLimitError,
    get_provider,
)
from g_gremlin.sequencer.models import BatchOutcome, Prospect

from ..api.progress import get_redis
from ..db import AsyncSessionLocal
from ..services.outreach_oauth import (
    OutreachOAuthError,
    refresh_if_needed,
    sanitize_outreach_error_text,
)
from ..services.secret_storage import get_saas_secret_backend
from .sequencer_row_contract import (
    build_submission_result,
    infer_provider_failure_reason,
    map_provider_outcome,
)
from .sequencer_worker_utils import (
    clean_string as _clean_string,
    get_row_number as _get_row_number,
    normalize_status as _normalize_status,
    should_skip_row as _should_skip_row,
)

logger = logging.getLogger(__name__)

STATE_KEY = "sequencer:job:{job_id}:state"
PROGRESS_KEY = "sequencer:job:{job_id}:progress"
SUMMARY_KEY = "sequencer:job:{job_id}:summary"
ERRORS_KEY = "sequencer:job:{job_id}:errors"
RESULTS_KEY = "sequencer:job:{job_id}:results"

RESULT_TTL = 86400
MAX_ERROR_COUNT = 50
MAX_RETRIES = 3
EMAIL_RE = re.compile(r"^[^@\s]+@[^@\s]+\.[^@\s]+$")

PROVIDER_CONFIG = {
    "smartlead": {
        "min_batch_delay": float(os.getenv("SMARTLEAD_MIN_BATCH_DELAY_SEC", "1.0")),
        "max_batch_size": int(os.getenv("SMARTLEAD_MAX_BATCH_SIZE", "350")),
        "required_fields": ["email"],
    },
    "instantly": {
        "min_batch_delay": float(os.getenv("INSTANTLY_MIN_BATCH_DELAY_SEC", "1.0")),
        "max_batch_size": int(os.getenv("INSTANTLY_MAX_BATCH_SIZE", "100")),
        "required_fields": ["email"],
    },
    "apollo": {
        "min_batch_delay": float(os.getenv("APOLLO_MIN_BATCH_DELAY_SEC", "0.5")),
        "max_batch_size": int(os.getenv("APOLLO_MAX_BATCH_SIZE", "1")),  # Apollo enrolls one at a time
        "required_fields": ["email"],
    },
    "outreach": {
        "min_batch_delay": float(os.getenv("OUTREACH_MIN_BATCH_DELAY_SEC", "1.0")),
        "max_batch_size": int(os.getenv("OUTREACH_MAX_BATCH_SIZE", "10")),
        "required_fields": ["email"],
    },
}
PROVIDER_LABELS = {
    "smartlead": "Smartlead",
    "instantly": "Instantly",
    "apollo": "Apollo",
    "outreach": "Outreach",
}

def _normalize_provider(value: Any) -> str:
    normalized = _clean_string(value).lower()
    return normalized or "smartlead"


def _provider_label(name: str) -> str:
    return PROVIDER_LABELS.get(name, name)


def _provider_secret_key_name(name: str) -> str:
    return "access_token" if name in {"apollo", "outreach"} else "api_key"


def _provider_uses_bearer_auth(name: str) -> bool:
    return name in {"apollo", "outreach"}


def _provider_requires_sender_account(name: str) -> bool:
    return name in {"apollo", "outreach"}

@dataclass
class RowContext:
    row_number: int
    email: str
    prospect: Prospect


class SequencerExportWorker:
    def __init__(
        self,
        account_id: str,
        job_id: str,
        config: Dict[str, Any],
        rows: List[Dict[str, Any]],
        redis=None,
        provider_factory=None,
        api_key: Optional[str] = None,
        sleep: Optional[callable] = None,
    ) -> None:
        self.account_id = account_id
        self.job_id = job_id
        self.config = config
        self.rows = rows
        self.redis = redis
        self.provider_factory = provider_factory
        self.api_key_override = api_key
        self.sleep = sleep or asyncio.sleep
        self.summary = {
            "total_rows": len(rows),
            "processed_rows": 0,
            "enrolled": 0,
            "exists": 0,
            "skipped": 0,
            "failed": 0,
        }
        self.errors: List[Dict[str, Any]] = []
        self.errors_truncated = False

    async def _get_redis(self):
        if self.redis is not None:
            return self.redis
        return await get_redis()

    async def _set_state(self, state: str) -> None:
        redis = await self._get_redis()
        if not redis:
            return
        await redis.set(STATE_KEY.format(job_id=self.job_id), state, ex=RESULT_TTL)

    async def _set_progress(self, progress: int) -> None:
        redis = await self._get_redis()
        if not redis:
            return
        await redis.set(PROGRESS_KEY.format(job_id=self.job_id), str(progress), ex=RESULT_TTL)

    async def _set_summary(self) -> None:
        redis = await self._get_redis()
        if not redis:
            return
        await redis.set(
            SUMMARY_KEY.format(job_id=self.job_id),
            json.dumps(self.summary),
            ex=RESULT_TTL,
        )

    async def _set_errors(self) -> None:
        redis = await self._get_redis()
        if not redis:
            return
        payload = {
            "errors": self.errors,
            "errors_truncated": self.errors_truncated,
        }
        await redis.set(
            ERRORS_KEY.format(job_id=self.job_id),
            json.dumps(payload),
            ex=RESULT_TTL,
        )

    async def _store_result(self, result: Dict[str, Any]) -> None:
        redis = await self._get_redis()
        if not redis:
            return
        await redis.rpush(RESULTS_KEY.format(job_id=self.job_id), json.dumps(result))
        await redis.expire(RESULTS_KEY.format(job_id=self.job_id), RESULT_TTL)

    def _record_error(self, row_number: int, status: str, error: str) -> None:
        if len(self.errors) >= MAX_ERROR_COUNT:
            self.errors_truncated = True
            return
        self.errors.append({
            "row_number": row_number,
            "status": status,
            "error": error,
        })

    def _increment_summary(self, status: str) -> None:
        if status in self.summary:
            self.summary[status] += 1
        self.summary["processed_rows"] = (
            self.summary["enrolled"]
            + self.summary["exists"]
            + self.summary["skipped"]
            + self.summary["failed"]
        )

    def _build_prospect(self, row: Dict[str, Any], mapping: Dict[str, Any]) -> Prospect:
        def value_for(field: str) -> Optional[str]:
            column = mapping.get(field)
            if not column:
                return None
            value = row.get(column)
            text = _clean_string(value)
            return text if text else None

        email = value_for("email") or ""
        custom_vars = {}
        for key, column in (mapping.get("custom_vars") or {}).items():
            if not column:
                continue
            raw = row.get(column)
            text = _clean_string(raw)
            if text:
                custom_vars[str(key)] = text

        return Prospect(
            email=email,
            first_name=value_for("first_name"),
            last_name=value_for("last_name"),
            company=value_for("company"),
            title=value_for("title"),
            phone=value_for("phone"),
            linkedin_url=value_for("linkedin_url"),
            domain=value_for("domain"),
            custom_vars=custom_vars,
            send_email_from_email_account_id=value_for("send_email_from_email_account_id"),
        )

    def _validate_row(
        self,
        row: Dict[str, Any],
        mapping: Dict[str, Any],
        required_fields: List[str],
    ) -> Optional[tuple[str, str]]:
        if "email" in required_fields:
            email_column = mapping.get("email")
            if not email_column:
                return "missing_required_field", "Email mapping is required"
            email_value = _clean_string(row.get(email_column))
            if not email_value:
                return "missing_required_field", "Missing required field"
            if not EMAIL_RE.match(email_value):
                return "invalid_email_format", "Invalid email format"
        return None

    async def _resolve_api_key(self) -> str:
        if self.api_key_override:
            return self.api_key_override
        profile_id = self.config.get("profile_id")
        provider_name = _normalize_provider(self.config.get("provider"))
        provider_label = _provider_label(provider_name)
        async with AsyncSessionLocal() as db:
            backend = get_saas_secret_backend()
            if provider_name == "outreach":
                try:
                    await refresh_if_needed(
                        db,
                        account_id=self.account_id,
                        force=False,
                        created_by="system",
                    )
                    await db.commit()
                except OutreachOAuthError as exc:
                    await db.rollback()
                    message = sanitize_outreach_error_text(str(exc))
                    if exc.status_code in {401, 403, 404}:
                        raise SequencerAuthError(message) from exc
                    raise SequencerError(message) from exc
            if profile_id:
                secret = await backend.get_secret_by_id(
                    db=db,
                    secret_id=profile_id,
                    workspace_id=self.account_id,
                    integration=provider_name,
                )
                if not secret or not secret.value:
                    raise SequencerAuthError("Profile not found")
                return secret.value
            key_name = _provider_secret_key_name(provider_name)
            secret = await backend.get_secret(
                db=db,
                integration=provider_name,
                key=key_name,
                workspace_id=self.account_id,
            )
        if not secret or not secret.value:
            raise SequencerAuthError(f"{provider_label} credentials not configured")
        return secret.value

    def _create_provider(self, api_key: str, provider_name: str):
        if self.provider_factory is None:
            kwargs = {}
            if _provider_uses_bearer_auth(provider_name):
                kwargs["use_bearer_auth"] = True
                send_from_id = _clean_string(self.config.get("send_email_from_email_account_id"))
                if send_from_id:
                    kwargs["send_email_from_email_account_id"] = send_from_id
            return get_provider(provider_name, api_key, **kwargs)
        try:
            return self.provider_factory(api_key=api_key)
        except TypeError:
            return self.provider_factory

    async def _push_batch_with_backoff(
        self,
        provider: Any,
        campaign_id: str,
        prospects: List[Prospect],
    ) -> BatchOutcome:
        attempt = 0
        logger.info(
            "[SEQUENCER-WORKER] _push_batch_with_backoff called: "
            "provider=%s campaign_id=%s prospect_count=%d emails=%s",
            getattr(provider, "name", "unknown"),
            campaign_id,
            len(prospects),
            [p.email for p in prospects[:3]],  # Log first 3 emails
        )
        while True:
            try:
                logger.info("[SEQUENCER-WORKER] Calling provider.upsert_prospects...")
                result = await asyncio.to_thread(
                    provider.upsert_prospects,
                    prospects,
                    campaign_id,
                )
                logger.info(
                    "[SEQUENCER-WORKER] upsert_prospects returned: succeeded=%d failed=%d skipped=%d",
                    len(result.succeeded),
                    len(result.failed),
                    len(result.skipped),
                )
                return result
            except SequencerRateLimitError as exc:
                retry_after = getattr(exc, "retry_after", None)
                delay = retry_after if retry_after is not None else (2 ** attempt)
                if attempt >= MAX_RETRIES:
                    raise
                await self.sleep(delay)
                attempt += 1
            except SequencerError:
                attempt += 1
                if attempt >= MAX_RETRIES:
                    raise
                await self.sleep(2 ** (attempt - 1))

    async def _apply_outcome(
        self,
        outcome: BatchOutcome,
        row_contexts: List[RowContext],
        campaign_id: str,
        provider_name: str,
    ) -> List[Dict[str, Any]]:
        rows_by_email: Dict[str, List[RowContext]] = {}
        for ctx in row_contexts:
            rows_by_email.setdefault(ctx.email.lower(), []).append(ctx)

        results: List[Dict[str, Any]] = []
        timestamp = datetime.now(timezone.utc).isoformat()

        def consume_row(email: str) -> Optional[RowContext]:
            queue = rows_by_email.get(email.lower())
            if not queue:
                return None
            return queue.pop(0)

        for group in (outcome.succeeded, outcome.failed, outcome.skipped):
            for item in group:
                ctx = consume_row(item.email)
                if not ctx:
                    continue
                submission_status, reason_codes = map_provider_outcome(
                    provider_name,
                    item.status,
                    item.error,
                )
                pushed_at = timestamp if submission_status in {"enrolled", "exists", "failed"} else ""
                submission_error = item.error if submission_status in {"failed", "skipped"} else ""
                result = build_submission_result(
                    row_number=ctx.row_number,
                    campaign_id=campaign_id,
                    submission_status=submission_status,
                    reason_codes=reason_codes,
                    provider_lead_id=item.provider_prospect_id,
                    pushed_at=pushed_at,
                    submission_error=submission_error,
                )
                results.append(result)

        for remaining in rows_by_email.values():
            for ctx in remaining:
                results.append(
                    build_submission_result(
                        row_number=ctx.row_number,
                        campaign_id=campaign_id,
                        submission_status="failed",
                        reason_codes=["batch_aborted"],
                        pushed_at="",
                        submission_error="No result returned",
                    )
                )

        return results

    async def _store_terminal_failures(
        self,
        *,
        row_numbers: List[int],
        campaign_id: str,
        reason_code: str,
        error_message: str,
        pushed_at: str = "",
        include_error_log: bool = True,
    ) -> None:
        for row_number in row_numbers:
            result = build_submission_result(
                row_number=row_number,
                campaign_id=campaign_id,
                submission_status="failed",
                reason_codes=[reason_code],
                pushed_at=pushed_at,
                submission_error=error_message,
            )
            await self._store_result(result)
            self._increment_summary("failed")
            if include_error_log:
                self._record_error(row_number, "failed", error_message)

    async def run(self) -> None:
        await self._set_state("running")
        await self._set_progress(0)

        provider_name = _normalize_provider(self.config.get("provider"))
        provider_config = PROVIDER_CONFIG.get(provider_name, PROVIDER_CONFIG["smartlead"])
        required_fields = provider_config.get("required_fields", ["email"])
        campaign_id = _clean_string(self.config.get("campaign_id"))
        force_reenroll = bool(self.config.get("force_reenroll"))
        mapping = self.config.get("mapping") or {}
        writeback = self.config.get("writeback") or {}
        all_row_numbers = [_get_row_number(row, idx) for idx, row in enumerate(self.rows)]

        if not campaign_id:
            await self._store_terminal_failures(
                row_numbers=all_row_numbers,
                campaign_id=campaign_id,
                reason_code="provider_validation_error",
                error_message="Missing campaign_id",
                pushed_at="",
            )
            await self._set_summary()
            await self._set_errors()
            await self._set_state("failed")
            return

        try:
            api_key = await self._resolve_api_key()
        except SequencerAuthError as exc:
            failure_timestamp = datetime.now(timezone.utc).isoformat()
            await self._store_terminal_failures(
                row_numbers=all_row_numbers,
                campaign_id=campaign_id,
                reason_code="provider_auth_error",
                error_message=str(exc),
                pushed_at=failure_timestamp,
            )
            await self._set_summary()
            await self._set_state("failed")
            await self._set_errors()
            return

        provider = self._create_provider(api_key, provider_name)

        pending_rows: List[RowContext] = []
        for idx, row in enumerate(self.rows):
            row_number = _get_row_number(row, idx)
            if not force_reenroll and _should_skip_row(row, campaign_id, writeback):
                result = build_submission_result(
                    row_number=row_number,
                    campaign_id=campaign_id,
                    submission_status="skipped",
                    reason_codes=["previously_enrolled"],
                    pushed_at="",
                    submission_error="Row already enrolled in this campaign",
                )
                await self._store_result(result)
                self._increment_summary("skipped")
                continue

            validation = self._validate_row(row, mapping, required_fields)
            if not validation and _provider_requires_sender_account(provider_name):
                default_send_from = _clean_string(self.config.get("send_email_from_email_account_id"))
                if not default_send_from:
                    sender_column = mapping.get("send_email_from_email_account_id")
                    sender_value = row.get(sender_column) if sender_column else None
                    if not _clean_string(sender_value):
                        message = (
                            "Missing mailbox id"
                            if provider_name == "outreach"
                            else "Missing send_email_from_email_account_id"
                        )
                        validation = ("no_sender_assigned", message)
            if validation:
                reason_code, message = validation
                result = build_submission_result(
                    row_number=row_number,
                    campaign_id=campaign_id,
                    submission_status="skipped",
                    reason_codes=[reason_code],
                    pushed_at="",
                    submission_error=message,
                )
                await self._store_result(result)
                self._increment_summary("skipped")
                continue

            prospect = self._build_prospect(row, mapping)
            email = _clean_string(prospect.email).lower()
            pending_rows.append(RowContext(row_number=row_number, email=email, prospect=prospect))

        await self._set_summary()
        await self._set_errors()

        total = self.summary["total_rows"]
        batch_size = max(1, min(int(provider_config.get("max_batch_size", 1)), 1000))
        min_delay = max(0.0, float(provider_config.get("min_batch_delay", 0.0)))
        next_index = 0

        try:
            for start in range(0, len(pending_rows), batch_size):
                next_index = start
                batch = pending_rows[start:start + batch_size]
                prospects = [ctx.prospect for ctx in batch]
                if min_delay:
                    await self.sleep(min_delay)
                outcome = await self._push_batch_with_backoff(provider, campaign_id, prospects)
                batch_results = await self._apply_outcome(outcome, batch, campaign_id, provider_name)
                for result in batch_results:
                    await self._store_result(result)
                    status = _normalize_status(
                        result.get("submission_status") or result.get("status") or "failed"
                    )
                    if status == "enrolled":
                        self._increment_summary("enrolled")
                    elif status == "exists":
                        self._increment_summary("exists")
                    elif status == "skipped":
                        self._increment_summary("skipped")
                    else:
                        self._increment_summary("failed")
                        error_message = result.get("submission_error") or result.get("error") or "Unknown error"
                        self._record_error(result.get("row_number", 0), "failed", error_message)

                next_index = start + len(batch)
                await self._set_summary()
                await self._set_errors()

                processed = self.summary["processed_rows"]
                progress = int(min(99, (processed / total) * 100)) if total else 100
                await self._set_progress(progress)

            await self._set_progress(100)
            await self._set_state("completed")
        except SequencerAuthError as exc:
            logger.error("Sequencer auth failure: %s", exc)
            remaining = pending_rows[next_index:]
            failure_timestamp = datetime.now(timezone.utc).isoformat()
            await self._store_terminal_failures(
                row_numbers=[ctx.row_number for ctx in remaining],
                campaign_id=campaign_id,
                reason_code="provider_auth_error",
                error_message=str(exc),
                pushed_at=failure_timestamp,
            )
            await self._set_summary()
            await self._set_errors()
            await self._set_state("failed")
        except Exception as exc:  # noqa: BLE001
            logger.exception("Sequencer export failed: %s", exc)
            remaining = pending_rows[next_index:]
            reason_code = infer_provider_failure_reason(str(exc))
            await self._store_terminal_failures(
                row_numbers=[ctx.row_number for ctx in remaining],
                campaign_id=campaign_id,
                reason_code=reason_code,
                error_message=str(exc),
                pushed_at=datetime.now(timezone.utc).isoformat(),
            )
            await self._set_summary()
            await self._set_errors()
            await self._set_state("failed")
        finally:
            if provider and hasattr(provider, "close"):
                try:
                    provider.close()
                except Exception:  # noqa: BLE001
                    logger.debug("Failed to close sequencer provider", exc_info=True)
