"""Async worker for pushing contacts to HeyReach campaigns."""

from __future__ import annotations

import asyncio
import csv
import io
import json
import logging
import random
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from g_gremlin.heyreach import (
    HeyReachClient,
    HeyReachProspect,
    HeyReachAuthError,
    HeyReachError,
    build_idempotency_key,
)

from ..api.progress import get_redis
from ..db import AsyncSessionLocal
from ..services.activation_utils import clean_string, normalize_status, get_row_number, should_skip_row
from ..services.email_notifications import send_heyreach_complete_notification
from ..services.heyreach_models import HeyReachConfig
from ..services.secret_storage import get_saas_secret_backend
from ..storage import get_storage_backend

logger = logging.getLogger(__name__)

META_KEY = "heyreach:job:{job_id}:meta"
STATE_KEY = "heyreach:job:{job_id}:state"
PROGRESS_KEY = "heyreach:job:{job_id}:progress"
SUMMARY_KEY = "heyreach:job:{job_id}:summary"
ERRORS_KEY = "heyreach:job:{job_id}:errors"
RESULTS_KEY = "heyreach:job:{job_id}:results"
CURSOR_KEY = "heyreach:job:{job_id}:cursor"

RESULT_TTL = 86400
MAX_ERROR_COUNT = 50

# Aliases for backward compatibility with internal references
_clean_string = clean_string
_normalize_status = normalize_status
_get_row_number = get_row_number
_should_skip_row = should_skip_row


@dataclass
class RowContext:
    row_number: int
    prospect: HeyReachProspect


class HeyReachWorker:
    """Process HeyReach push job asynchronously."""

    CHUNK_SIZE = 10

    def __init__(
        self,
        account_id: str,
        job_id: str,
        config: Dict[str, Any],
        rows: Optional[List[Dict[str, Any]]] = None,
        redis=None,
        sleep: Optional[callable] = None,
    ) -> None:
        self.account_id = account_id
        self.job_id = job_id
        self.config = HeyReachConfig.model_validate(config).model_dump()
        self.rows = rows or []
        self.redis = redis
        self.sleep = sleep or asyncio.sleep
        self.summary = {
            "total_rows": len(self.rows),
            "processed_rows": 0,
            "enrolled": 0,
            "failed": 0,
            "skipped": 0,
            "invalid": 0,
            "duplicate": 0,
        }
        self.errors: List[Dict[str, Any]] = []
        self.errors_truncated = False
        self._resume_cursor = 0

    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, message: Optional[str] = None) -> 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)
        await self._update_job_meta(status=state, message=message)

    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)
        await self._update_job_meta(progress=progress, processed_rows=self.summary["processed_rows"])

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

    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)
        await self._update_job_meta(errors=payload)

    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)

    async def _update_job_meta(self, **updates: Any) -> None:
        redis = await self._get_redis()
        if not redis:
            return
        key = f"job:{self.job_id}"
        raw = await redis.get(key)
        meta: Dict[str, Any]
        if raw:
            try:
                meta = json.loads(raw)
            except Exception:
                meta = {}
        else:
            meta = {}
        meta.update({k: v for k, v in updates.items() if v is not None})
        await redis.set(key, json.dumps(meta), ex=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["failed"]
            + self.summary["skipped"]
            + self.summary["invalid"]
            + self.summary["duplicate"]
        )

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

    async def _load_rows(self) -> None:
        if self.rows:
            self.summary["total_rows"] = len(self.rows)
            return

        upload_id = self.config.get("upload_id")
        if not upload_id:
            return

        storage = get_storage_backend()
        if asyncio.iscoroutine(storage):
            storage = await storage
        content = await storage.read(upload_id)
        text = content.decode("utf-8-sig", errors="ignore")
        reader = csv.DictReader(io.StringIO(text))
        rows = list(reader)
        starting_row = int(self.config.get("starting_row") or 2)
        for idx, row in enumerate(rows):
            row["row_number"] = starting_row + idx
        self.rows = rows
        self.summary["total_rows"] = len(rows)

    async def _get_client(self) -> HeyReachClient:
        profile_id = self.config.get("profile_id")
        if not profile_id:
            raise HeyReachAuthError("Profile not configured")
        async with AsyncSessionLocal() as db:
            backend = get_saas_secret_backend()
            secret = await backend.get_secret_by_id(
                db=db,
                secret_id=profile_id,
                workspace_id=self.account_id,
                integration="heyreach",
            )
            workspace_secret = await backend.get_secret(
                db=db,
                integration="heyreach",
                key="workspace_id",
                workspace_id=self.account_id,
            )
        if not secret or not secret.value:
            raise HeyReachAuthError("HeyReach credentials not configured")
        workspace_id = None
        if workspace_secret and workspace_secret.value:
            workspace_id = _clean_string(workspace_secret.value)
        return HeyReachClient(api_key=str(secret.value), workspace_id=workspace_id)

    def _build_prospect(self, row: Dict[str, Any], mapping: Dict[str, Any]) -> Optional[HeyReachProspect]:
        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")
        linkedin_url = value_for("linkedin_url")
        if not email or not linkedin_url:
            return None

        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 HeyReachProspect(
            email=email,
            linkedin_url=linkedin_url,
            first_name=value_for("first_name"),
            last_name=value_for("last_name"),
            company=value_for("company"),
            title=value_for("title"),
            sender_account_id=value_for("sender_account_id"),
            custom_vars=custom_vars,
        )

    def _build_batch_results(
        self,
        outcome,
        row_contexts: List[RowContext],
        campaign_id: str,
    ) -> List[Dict[str, Any]]:
        contexts_by_email: Dict[str, List[RowContext]] = {}
        for ctx in row_contexts:
            key = ctx.prospect.email.strip().lower()
            contexts_by_email.setdefault(key, []).append(ctx)

        def pop_ctx(email: str) -> Optional[RowContext]:
            key = (email or "").strip().lower()
            if not key:
                return None
            bucket = contexts_by_email.get(key)
            if not bucket:
                return None
            return bucket.pop(0)

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

        for item in outcome.succeeded or []:
            ctx = pop_ctx(item.get("prospect_email", ""))
            if not ctx:
                continue
            status = _clean_string(item.get("status") or "enrolled")
            if status == "duplicate":
                self._increment_summary("duplicate")
            else:
                self._increment_summary("enrolled")
            results.append({
                "row_number": ctx.row_number,
                "status": status,
                "campaign_id": campaign_id,
                "provider_prospect_id": item.get("provider_prospect_id"),
                "pushed_at": now_iso,
            })

        for item in (outcome.failed or []) + (outcome.retryable_failed or []):
            ctx = pop_ctx(item.get("prospect_email", ""))
            if not ctx:
                continue
            error = _clean_string(item.get("error") or "Failed to enroll")
            self._record_error(ctx.row_number, "failed", error)
            self._increment_summary("failed")
            results.append({
                "row_number": ctx.row_number,
                "status": "failed",
                "campaign_id": campaign_id,
                "error": error,
                "pushed_at": now_iso,
            })

        for remaining in contexts_by_email.values():
            for ctx in remaining:
                error = "No enrollment result returned"
                self._record_error(ctx.row_number, "failed", error)
                self._increment_summary("failed")
                results.append({
                    "row_number": ctx.row_number,
                    "status": "failed",
                    "campaign_id": campaign_id,
                    "error": error,
                    "pushed_at": now_iso,
                })

        return results

    async def run(self) -> None:
        adapter: Optional[HeyReachClient] = None
        try:
            await self._set_state("processing", message="Starting HeyReach push")
            await self._load_rows()

            if not self.rows:
                await self._set_summary()
                await self._set_state("completed", message="No rows to process")
                await self._set_progress(100)
                return

            adapter = await self._get_client()
            campaign_id = self.config.get("campaign_id")
            mapping = self.config.get("mapping") or {}
            writeback = self.config.get("writeback") or {}
            starting_row = int(self.config.get("starting_row") or 2)
            total_rows = int(self.summary.get("total_rows") or len(self.rows))

            row_contexts: List[RowContext] = []
            seen_emails: set[str] = set()

            for idx, row in enumerate(self.rows):
                row_number = _get_row_number(row, idx, starting_row)

                if _should_skip_row(row, campaign_id, writeback):
                    await self._store_result({
                        "row_number": row_number,
                        "status": "skipped",
                        "campaign_id": campaign_id,
                        "error": "Already enrolled",
                    })
                    self._increment_summary("skipped")
                    await self._set_resume_cursor(idx + 1)
                    continue

                prospect = self._build_prospect(row, mapping)
                if not prospect:
                    await self._store_result({
                        "row_number": row_number,
                        "status": "invalid",
                        "campaign_id": campaign_id,
                        "error": "Missing required email or LinkedIn URL",
                    })
                    self._record_error(row_number, "invalid", "Missing required email or LinkedIn URL")
                    self._increment_summary("invalid")
                    await self._set_resume_cursor(idx + 1)
                    continue

                email_key = prospect.email.strip().lower()
                if email_key in seen_emails:
                    await self._store_result({
                        "row_number": row_number,
                        "status": "duplicate",
                        "campaign_id": campaign_id,
                        "error": "Duplicate email in selection",
                    })
                    self._increment_summary("duplicate")
                    await self._set_resume_cursor(idx + 1)
                    continue

                seen_emails.add(email_key)
                row_contexts.append(RowContext(row_number=row_number, prospect=prospect))

                if len(row_contexts) >= self.CHUNK_SIZE:
                    batch_prospects = [ctx.prospect for ctx in row_contexts]
                    idempotency_keys = [
                        build_idempotency_key(campaign_id, ctx.prospect.email)
                        for ctx in row_contexts
                    ]
                    outcome = adapter.enroll_prospects(
                        batch_prospects,
                        campaign_id,
                        idempotency_keys=idempotency_keys,
                    )
                    payloads = self._build_batch_results(outcome, row_contexts, campaign_id)
                    for payload in payloads:
                        await self._store_result(payload)
                    await self._set_summary()
                    await self._set_errors()
                    progress = int((self.summary["processed_rows"] / max(total_rows, 1)) * 100)
                    await self._set_progress(min(progress, 99))
                    row_contexts = []
                    await self._set_resume_cursor(idx + 1)
                    await self.sleep(0.5 * random.uniform(0.5, 1.5))

            if row_contexts:
                batch_prospects = [ctx.prospect for ctx in row_contexts]
                idempotency_keys = [
                    build_idempotency_key(campaign_id, ctx.prospect.email)
                    for ctx in row_contexts
                ]
                outcome = adapter.enroll_prospects(
                    batch_prospects,
                    campaign_id,
                    idempotency_keys=idempotency_keys,
                )
                payloads = self._build_batch_results(outcome, row_contexts, campaign_id)
                for payload in payloads:
                    await self._store_result(payload)

            await self._set_summary()
            await self._set_errors()
            await self._set_progress(100)
            await self._set_state("completed", message="HeyReach push completed")
            await self._send_completion_notification()
        except (HeyReachAuthError, HeyReachError) as exc:
            await self._set_state("failed", message=str(exc))
            await self._set_progress(100)
            raise
        except Exception as exc:
            await self._set_state("failed", message=str(exc))
            await self._set_progress(100)
            raise

    async def _send_completion_notification(self) -> None:
        """Send email notification when job completes."""
        notify_email = self.config.get("notify_email")
        if not notify_email:
            return
        spreadsheet_url = self.config.get("spreadsheet_url")
        try:
            await send_heyreach_complete_notification(
                user_email=notify_email,
                job_name="HeyReach LinkedIn Push",
                job_id=self.job_id,
                total_rows=self.summary.get("total_rows", 0),
                enrolled_count=self.summary.get("enrolled", 0),
                failed_count=self.summary.get("failed", 0),
                skipped_count=self.summary.get("skipped", 0),
                duplicate_count=self.summary.get("duplicate", 0),
                invalid_count=self.summary.get("invalid", 0),
                spreadsheet_url=spreadsheet_url,
            )
        except Exception as exc:
            logger.warning(f"Failed to send completion notification: {exc}")
