import asyncio
import random
import time
from typing import Any, Optional

from loguru import logger

from app.core.clickhouse import clickhouse_manager
from app.models.company import CompanyCleaningQueue
from app.models.token import BossToken
from app.services.company_jobs_sync import CompanyJobsSyncService
from app.services.company_storage import company_storage, normalize_company_id
from app.services.crawler.boss import BossService
from app.services.crawler.qcwy import QcwyService
from app.services.crawler.zhilian import ZhilianService


SOURCE_CONFIGS: dict[str, dict[str, Any]] = {
    "zhilian": {
        "job_table": "zhilian_job",
        "company_id_expr": "JSONExtractString(json_data, 'companyNumber')",
        "company_name_expr": "JSONExtractString(json_data, 'companyName')",
        "days_back": 30,
        "max_query_limit": None,
    },
    "qcwy": {
        "job_table": "qcwy_job",
        "company_id_expr": "JSONExtractString(json_data, 'coId')",
        "company_name_expr": "JSONExtractString(json_data, 'companyName')",
        "days_back": 30,
        "max_query_limit": 5000,
    },
    "boss": {
        "job_table": "boss_job",
        "company_id_expr": "JSONExtractString(json_data, 'brandComInfoVO', 'encryptBrandId')",
        "company_name_expr": "JSONExtractString(json_data, 'brandComInfoVO', 'brandName')",
        "days_back": 30,
        "max_query_limit": None,
    },
}


class CompanyCleaner:
    _TOKEN_REFRESH_INTERVAL = 3600

    def __init__(self):
        self.boss_service = BossService()
        self.qcwy_service = QcwyService()
        self.zhilian_service = ZhilianService()
        self.company_jobs_sync = CompanyJobsSyncService()
        self._boss_token_loaded = False
        self._token_loaded_at: float = 0

    def _apply_proxy(self, proxy: Optional[str]) -> None:
        self.boss_service.set_proxy(proxy)
        self.qcwy_service.set_proxy(proxy)
        self.zhilian_service.set_proxy(proxy)
        self.company_jobs_sync.set_proxy(proxy)

    async def _ensure_boss_token_loaded(self) -> None:
        now = time.time()
        if (
            self._boss_token_loaded
            and self.boss_service.login_data.get("mpt")
            and now - self._token_loaded_at < self._TOKEN_REFRESH_INTERVAL
        ):
            return
        token_obj = await BossToken.filter(is_active=True).order_by("-updated_at").first()
        if not token_obj:
            logger.warning("BossToken not found or inactive in CompanyCleaner")
            return
        self.boss_service.set_login_data(token_obj.mpt or "", "")
        self._boss_token_loaded = True
        self._token_loaded_at = now

    async def collect_pending_companies(self, limit: int = 1000, source: Optional[str] = None) -> dict[str, Any]:
        client = await clickhouse_manager.get_client()
        logger.info(f"Starting to collect pending companies (limit={limit}, source={source or 'all'})...")
        summary: dict[str, Any] = {
            "total_created": 0,
            "sources": {},
        }
        sources = [s for s in SOURCE_CONFIGS if source is None or source == s]
        # 并行采集各平台
        tasks = [self._collect_source(client, s, limit) for s in sources]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        for s, result in zip(sources, results):
            if isinstance(result, Exception):
                logger.error(f"Error collecting {s}: {result}")
                summary["sources"][s] = {"source": s, "created_count": 0, "error": str(result)}
            else:
                summary["sources"][s] = result
                summary["total_created"] += result["created_count"]
        logger.info("Finished collecting pending companies.")
        return summary

    async def _collect_source(self, client, source: str, limit: int) -> dict[str, Any]:
        config = SOURCE_CONFIGS[source]

        # 先从 MySQL 取出该平台所有已入队/已入库的 company_id，Python 侧快速排除
        all_queued = set(await CompanyCleaningQueue.filter(source=source).values_list("company_id", flat=True))
        all_existing = await company_storage.get_all_company_ids(source)
        exclude_ids = all_queued | all_existing
        logger.info(f"Loaded {len(exclude_ids)} known {source} company IDs for exclusion")

        # 用 OFFSET 跳过已知公司数量，获取新公司
        offset = len(exclude_ids)
        result = await self._query_candidate_rows(
            client=client,
            table=config["job_table"],
            company_id_expr=config["company_id_expr"],
            company_name_expr=config["company_name_expr"],
            days_back=config["days_back"],
            limit=limit,
            max_query_limit=config["max_query_limit"],
            offset=offset,
        )
        if not result.result_rows:
            logger.info(f"No new {source} companies found in job table query.")
            return {
                "source": source,
                "query_count": 0,
                "deduped_count": 0,
                "existing_count": len(all_existing),
                "queued_count": len(all_queued),
                "created_count": 0,
            }

        deduped_candidates: list[dict[str, str]] = []
        seen_ids: set[str] = set()
        for raw_company_id, company_name in result.result_rows:
            company_id = normalize_company_id(source, raw_company_id)
            if not company_id or company_id in seen_ids or company_id in exclude_ids:
                continue
            seen_ids.add(company_id)
            deduped_candidates.append(
                {
                    "company_id": company_id,
                    "company_name": (company_name or "").strip(),
                }
            )
            if len(deduped_candidates) >= limit:
                break

        created_count = await company_storage.enqueue_companies(source, deduped_candidates)
        logger.info(f"Added {created_count} {source} companies to MySQL queue.")
        return {
            "source": source,
            "query_count": len(result.result_rows),
            "deduped_count": len(deduped_candidates),
            "existing_count": len(all_existing),
            "queued_count": len(all_queued),
            "created_count": created_count,
        }

    async def _query_candidate_rows(
        self,
        *,
        client,
        table: str,
        company_id_expr: str,
        company_name_expr: str,
        days_back: int,
        limit: int,
        max_query_limit: Optional[int],
        offset: int = 0,
    ):
        current_days = days_back
        current_limit = limit * 5
        if max_query_limit is not None:
            current_limit = min(current_limit, max_query_limit)

        last_error: Optional[Exception] = None
        for attempt in range(3):
            sample_sql = " SAMPLE 0.1" if attempt == 2 else ""
            current_offset = offset
            if attempt == 1:
                current_days = max(1, min(current_days, 3))
                current_limit = min(current_limit, max(limit, 50))
                current_offset = 0
            query = f"""
            SELECT DISTINCT
                {company_id_expr} AS cid,
                {company_name_expr} AS cname
            FROM job_data.{table}{sample_sql}
            PREWHERE created_at > now() - INTERVAL {current_days} DAY
                AND json_data != ''
            WHERE {company_id_expr} != ''
            LIMIT {current_limit} OFFSET {current_offset}
            """
            try:
                logger.info(
                    f"Querying company candidates from {table} "
                    f"(days={current_days}, limit={current_limit}, attempt={attempt + 1})"
                )
                return await client.query(query)
            except Exception as exc:
                last_error = exc
                error_str = str(exc).lower()
                if "memory" in error_str or "memory_limit" in error_str:
                    logger.warning(f"Memory-sensitive query retry for {table}: {exc}")
                    continue
                raise
        assert last_error is not None
        raise last_error

    async def process_single_company(
        self,
        source: str,
        company_id: str,
        proxy: Optional[str] = None,
        max_delay_seconds: int = 5,
    ) -> dict[str, Any]:
        normalized_id = normalize_company_id(source, company_id)
        queue, _ = await company_storage.enqueue_company(source, normalized_id)
        if proxy:
            self._apply_proxy(proxy)
        if max_delay_seconds and max_delay_seconds > 0:
            await asyncio.sleep(random.randint(1, max_delay_seconds))

        await company_storage.mark_queue_processing(queue)
        try:
            persist_result = await self._fetch_and_save(source, normalized_id)
            jobs_result = await self._sync_company_jobs(source, normalized_id)
            if persist_result["company_name"] and queue.company_name != persist_result["company_name"]:
                queue.company_name = persist_result["company_name"]
            await company_storage.mark_queue_result(queue, status="done", jobs_summary=jobs_result)
            return {
                "success": True,
                "source": source,
                "company_id": normalized_id,
                "company_name": persist_result["company_name"],
                "status": "done",
                "error_msg": "",
                "created": persist_result["created"],
                "jobs_summary": jobs_result,
            }
        except Exception as exc:
            logger.error(f"Error processing {source} {normalized_id}: {exc}")
            await company_storage.mark_queue_result(
                queue,
                status="failed",
                error_msg=str(exc),
                increment_retry=True,
            )
            return {
                "success": False,
                "source": source,
                "company_id": normalized_id,
                "company_name": queue.company_name or "",
                "status": "failed",
                "error_msg": str(exc),
            }

    async def process_pending_companies(
        self,
        limit: int = 100,
        source: Optional[str] = None,
        proxy: Optional[str] = None,
        max_delay_seconds: int = 0,
    ):
        logger.info(f"Processing pending companies (limit={limit}, source={source or 'all'})...")
        if proxy:
            self._apply_proxy(proxy)

        query = CompanyCleaningQueue.filter(status="pending")
        if source:
            query = query.filter(source=source)
        queue_rows = await query.order_by("created_at").limit(limit)
        if not queue_rows:
            logger.info("No pending companies to process.")
            return

        for queue in queue_rows:
            logger.info(f"Processing {queue.source} company: {queue.company_name or ''} ({queue.company_id})")
            await company_storage.mark_queue_processing(queue)
            try:
                if max_delay_seconds and max_delay_seconds > 0:
                    await asyncio.sleep(random.randint(1, max_delay_seconds))
                persist_result = await self._fetch_and_save(queue.source, queue.company_id)
                jobs_result = await self._sync_company_jobs(queue.source, queue.company_id)
                logger.info(
                    f"Synced {queue.source} company jobs: "
                    f"fetched={jobs_result['jobs_fetched']} stored={jobs_result['stored_success']} "
                    f"duplicate={jobs_result['duplicate']} failed={jobs_result['failed']}"
                )
                if persist_result["company_name"] and queue.company_name != persist_result["company_name"]:
                    queue.company_name = persist_result["company_name"]
                await company_storage.mark_queue_result(queue, status="done", jobs_summary=jobs_result)
            except Exception as exc:
                logger.error(f"Error processing {queue.source} {queue.company_id}: {exc}")
                await company_storage.mark_queue_result(
                    queue,
                    status="failed",
                    error_msg=str(exc),
                    increment_retry=True,
                )

    async def _fetch_and_save(self, source: str, company_id: str) -> dict[str, Any]:
        data = await self._fetch_company_data(source, company_id)
        if not data:
            raise ValueError(f"No data returned from source={source} company_id={company_id}")
        return await company_storage.upsert_company(source, data, company_id=company_id)

    async def _sync_company_jobs(self, source: str, company_id: str) -> dict[str, Any]:
        try:
            return await self.company_jobs_sync.sync_company_jobs(source, company_id)
        except Exception as exc:
            logger.warning(f"Sync company jobs failed for {source} {company_id}: {exc}")
            return {
                "success": False,
                "source": source,
                "company_id": company_id,
                "jobs_fetched": 0,
                "stored_success": 0,
                "duplicate": 0,
                "failed": 0,
                "error": str(exc),
            }

    async def _fetch_company_data(self, source: str, company_id: str) -> dict[str, Any]:
        if source == "zhilian":
            data = await asyncio.to_thread(self.zhilian_service.get_company_detail, company_id)
            return data or {}
        if source == "qcwy":
            data = await asyncio.to_thread(self.qcwy_service.get_company_info, company_id)
            return data or {}
        if source == "boss":
            await self._ensure_boss_token_loaded()
            data = await asyncio.to_thread(self.boss_service.get_company_detail_by_id, company_id)
            return data or {}
        raise ValueError(f"unsupported source: {source}")

    async def cleanup_old_records(self):
        logger.info("Starting cleanup of processed pending companies...")
        await CompanyCleaningQueue.filter(status__in=["done", "failed"]).delete()


company_cleaner = CompanyCleaner()
