"""Async worker for Company List Builder exports."""

from __future__ import annotations

import asyncio
import csv
import inspect
import json
import logging
import os
import shutil
import tempfile
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Dict, Iterable, List, Tuple, TYPE_CHECKING
from uuid import uuid4

# Lazy imports for GCP modules to avoid credential discovery hang on Windows
if TYPE_CHECKING:
    pass

_gcp_modules = {}

def _get_bigquery():
    """Lazily import BigQuery to avoid GCP credential discovery at startup."""
    if "bigquery" not in _gcp_modules:
        from google.cloud import bigquery
        _gcp_modules["bigquery"] = bigquery
    return _gcp_modules["bigquery"]

def _get_storage():
    """Lazily import Storage to avoid GCP credential discovery at startup."""
    if "storage" not in _gcp_modules:
        from google.cloud import storage
        _gcp_modules["storage"] = storage
    return _gcp_modules["storage"]

from ..db import AsyncSessionLocal
from ..services.company_list_service import CompanyQueryBuilder
from ..services.list_builder_quota import ListBuilderQuota
from .. import settings as saas_settings
from ..storage import LocalStorageBackend, get_storage_backend

logger = logging.getLogger(__name__)

STATE_KEY = "listbuilder:export:{job_id}:state"
PROGRESS_KEY = "listbuilder:export:{job_id}:progress"
RESULT_KEY = "listbuilder:export:{job_id}:result"
IDEMPOTENCY_KEY = "listbuilder:export:{account_id}:{key}"
RESULT_TTL = 86400  # 24h
IDEMPOTENCY_TTL = 86400  # 24h
ACTIVE_JOBS_TTL = 7200  # 2h default to avoid stale locks
PROGRESS_CAP = 95  # leave room for upload + finalization


def _make_absolute_url(url: str) -> str:
    if not url or url.startswith("http://") or url.startswith("https://"):
        return url
    base = os.getenv("PUBLIC_BASE_URL")
    if not base:
        if saas_settings.ENV.lower() in ("dev", "development"):
            base = os.getenv("FM_API_BASE_URL", "http://localhost:8000")
        else:
            base = getattr(saas_settings, "PUBLIC_BASE_URL", "") or ""
    if not base:
        return url
    base = str(base).rstrip("/")
    normalized = url if str(url).startswith("/") else f"/{url}"
    return f"{base}{normalized}"


class ListBuilderWorker:
    """Background exporter that streams BigQuery results to GCS."""

    def __init__(self, redis=None, bucket_name: str | None = None) -> None:
        self.redis = redis
        self.query_builder = CompanyQueryBuilder()
        env_bucket = (
            os.getenv("FM_LISTBUILDER_EXPORT_BUCKET")
            or os.getenv("FM_EXPORT_BUCKET")
            or os.getenv("FM_FOUNDRYGRAPH_BUCKET")
        )
        resolved_bucket = bucket_name or env_bucket
        self.bucket_name = resolved_bucket or "foundrymatch-exports"
        self._use_local_storage = (
            saas_settings.STORAGE_BACKEND == "local" and resolved_bucket is None
        )
        self._storage_client = None
        if not self._use_local_storage:
            storage_mod = _get_storage()
            self._storage_client = (
                storage_mod.Client(project=self.query_builder.PROJECT_ID)
                if self.query_builder.PROJECT_ID
                else storage_mod.Client()
            )

    def _active_jobs_key(self, account_id: str) -> str:
        return f"listbuilder:account:{account_id}:jobs"

    def _idempotency_cache_key(self, account_id: str, key: str) -> str:
        return IDEMPOTENCY_KEY.format(account_id=account_id, key=key)

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

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

    async def _set_result(self, job_id: str, payload: Dict[str, Any]) -> None:
        if not self.redis:
            return
        await self.redis.set(
            RESULT_KEY.format(job_id=job_id),
            json.dumps(payload),
            ex=RESULT_TTL,
        )

    async def start_export_job(
        self,
        account_id: str,
        filters: Dict[str, Any],
        columns: Iterable[str],
        limit: int,
        sort_by: str | None = None,
        sort_dir: str | None = None,
        idempotency_key: str | None = None,
    ) -> Tuple[str, bool]:
        """
        Start an export. Returns (job_id, started_new_job).
        If idempotency key matches, returns existing job id and False.
        """
        if self.redis and idempotency_key:
            cache_key = self._idempotency_cache_key(account_id, idempotency_key)
            existing = await self.redis.get(cache_key)
            if existing:
                return existing.decode("utf-8") if isinstance(existing, (bytes, bytearray)) else str(existing), False

        job_id = str(uuid4())

        if self.redis:
            await self.redis.sadd(self._active_jobs_key(account_id), job_id)
            await self.redis.expire(self._active_jobs_key(account_id), ACTIVE_JOBS_TTL)
            if idempotency_key:
                await self.redis.set(
                    self._idempotency_cache_key(account_id, idempotency_key),
                    job_id,
                    ex=IDEMPOTENCY_TTL,
                )
            await self._set_state(job_id, "queued")
            await self._set_progress(job_id, 0)

        asyncio.create_task(
            self._run_job(
                job_id=job_id,
                account_id=account_id,
                filters=filters,
                columns=columns,
                limit=limit,
                sort_by=sort_by,
                sort_dir=sort_dir,
            )
        )

        return job_id, True

    async def _run_job(
        self,
        job_id: str,
        account_id: str,
        filters: Dict[str, Any],
        columns: Iterable[str],
        limit: int,
        sort_by: str | None,
        sort_dir: str | None,
    ) -> None:
        await self._set_state(job_id, "processing")
        await self._set_progress(job_id, 5)

        query, params, selected_columns = self.query_builder.build_query(
            filters=filters,
            columns=columns,
            limit=limit,
            offset=0,
            sort_by=sort_by,
            sort_dir=sort_dir,
            include_extra_row=False,
        )

        try:
            row_count, tmp_path = await asyncio.to_thread(
                self._export_to_tempfile, query, params, selected_columns
            )
        except Exception as exc:  # noqa: BLE001
            logger.exception("List builder export failed during query: %s", exc)
            await self._set_state(job_id, "failed")
            await self._set_result(job_id, {"error": str(exc)})
            await self._cleanup_active(account_id, job_id)
            return

        try:
            await self._set_progress(job_id, min(PROGRESS_CAP, 90))
            signed_url, expires_at = await asyncio.to_thread(
                self._upload_and_presign, tmp_path, account_id, job_id
            )
        except Exception as exc:  # noqa: BLE001
            logger.exception("List builder export failed during upload: %s", exc)
            await self._set_state(job_id, "failed")
            await self._set_result(job_id, {"error": str(exc)})
            await self._cleanup_active(account_id, job_id)
            Path(tmp_path).unlink(missing_ok=True)
            return
        finally:
            Path(tmp_path).unlink(missing_ok=True)

        async with AsyncSessionLocal() as db:
            quota = ListBuilderQuota(db=db, redis=self.redis)
            try:
                await quota.record_export(
                    account_id,
                    row_count,
                    idempotency_key=f"list_builder:{job_id}",
                )
            except Exception as exc:  # noqa: BLE001
                logger.warning("Failed to record list builder usage for %s: %s", account_id, exc)
            await quota.clear_job(account_id, job_id)

        await self._set_result(
            job_id,
            {
                "download_url": signed_url,
                "rows": row_count,
                "expires_at": expires_at.isoformat(),
            },
        )
        await self._set_progress(job_id, 100)
        await self._set_state(job_id, "completed")

    def _export_to_tempfile(
        self,
        query: str,
        params: List[Any],
        columns: List[str],
    ) -> Tuple[int, str]:
        bq = _get_bigquery()
        job = self.query_builder.client.query(
            query, job_config=bq.QueryJobConfig(query_parameters=params)
        )
        result = job.result(page_size=1000)

        tmp = tempfile.NamedTemporaryFile(
            delete=False,
            suffix=".csv",
            dir=tempfile.gettempdir(),
            mode="w",
            newline="",
            encoding="utf-8",
        )
        writer = csv.writer(tmp)
        writer.writerow(columns)

        row_count = 0
        for row in result:
            writer.writerow([CompanyQueryBuilder._coerce(row[col]) for col in columns])
            row_count += 1

        tmp.flush()
        tmp.close()
        return row_count, tmp.name

    def _upload_and_presign(self, path: str, account_id: str, job_id: str) -> Tuple[str, datetime]:
        if self._use_local_storage:
            return self._upload_local(path, account_id, job_id)

        if not self._storage_client:
            raise RuntimeError("Storage client unavailable for list builder export.")

        bucket = self._storage_client.bucket(self.bucket_name)
        blob = bucket.blob(f"exports/{account_id}/{job_id}/companies.csv")
        try:
            blob.upload_from_filename(path, content_type="text/csv")
        except Exception as e:  # noqa: BLE001
            if saas_settings.ENV.lower() in ("dev", "development"):
                logger.warning(
                    "GCS upload failed for job %s; falling back to local storage. Error: %s",
                    job_id,
                    e,
                )
                return self._upload_local(path, account_id, job_id)
            raise

        expires_at = datetime.now(timezone.utc) + timedelta(hours=1)
        try:
            signed_url = blob.generate_signed_url(expiration=timedelta(hours=1), method="GET")
        except Exception as e:  # noqa: BLE001
            logger.error(
                "Signed URL generation failed for job %s. "
                "Ensure service account has iam.serviceAccounts.signBlob permission. Error: %s",
                job_id,
                e,
            )
            raise RuntimeError(
                "Signed URL generation failed. Ensure Cloud Run service account can sign URLs. "
                f"{e}"
            ) from e
        return signed_url, expires_at

    def _upload_local(self, path: str, account_id: str, job_id: str) -> Tuple[str, datetime]:
        storage = get_storage_backend()
        if not hasattr(storage, "base_path"):
            storage = LocalStorageBackend(saas_settings.LOCAL_STORAGE_PATH)
        key = f"exports/{account_id}/{job_id}/companies.csv"

        if hasattr(storage, "base_path"):
            dest_path = storage._resolve_path(key)
        else:
            dest_path = Path(saas_settings.LOCAL_STORAGE_PATH).resolve() / key

        dest_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.copyfile(path, dest_path)

        expires_at = datetime.now(timezone.utc) + timedelta(hours=1)
        try:
            url = storage.generate_presigned_url(
                key=key, expires_in=3600, operation="get_object"
            )
            if inspect.isawaitable(url):
                url = asyncio.run(url)
        except Exception as exc:  # noqa: BLE001
            logger.warning("Failed to generate local export URL for %s: %s", key, exc)
            url = str(dest_path)

        return _make_absolute_url(str(url)), expires_at

    async def _cleanup_active(self, account_id: str, job_id: str) -> None:
        if self.redis:
            await self.redis.srem(self._active_jobs_key(account_id), job_id)
