from __future__ import annotations

from datetime import datetime
from typing import Any, Dict, Iterable, Optional, Type

from app.models.company import (
    BaseCompanyModel,
    BossCompany,
    CompanyCleaningQueue,
    QcwyCompany,
    ZhilianCompany,
)


COMPANY_SOURCES = {"boss", "qcwy", "zhilian"}
QUEUE_TERMINAL_STATUSES = {"done", "failed"}


def normalize_company_id(source: str, company_id: str) -> str:
    value = str(company_id or "").strip()
    if source == "qcwy" and value.lower().startswith("co") and value[2:].isdigit():
        return value[2:]
    return value


def _pick_first(data: dict[str, Any], *keys: str) -> Optional[Any]:
    for key in keys:
        value = data.get(key)
        if value not in (None, ""):
            return value
    return None


def _nested_get(data: dict[str, Any], *path: str) -> Any:
    current: Any = data
    for key in path:
        if not isinstance(current, dict):
            return None
        current = current.get(key)
    return current


def _clean_text(value: Any) -> Optional[str]:
    if value is None:
        return None
    text = str(value).strip()
    return text or None


def _model_for_source(source: str) -> Type[BaseCompanyModel]:
    mapping: dict[str, Type[BaseCompanyModel]] = {
        "boss": BossCompany,
        "qcwy": QcwyCompany,
        "zhilian": ZhilianCompany,
    }
    if source not in mapping:
        raise ValueError(f"unsupported source: {source}")
    return mapping[source]


def _extract_boss_fields(raw: dict[str, Any], company_id: str) -> dict[str, Any]:
    payload = raw.get("zpData") if isinstance(raw.get("zpData"), dict) else raw
    brand = payload.get("brandComInfoVO") or {}
    company_full = payload.get("companyFullInfoVO") or {}

    return {
        "source_company_id": normalize_company_id("boss", company_id or _pick_first(brand, "encryptBrandId", "brandId")),
        "company_name": _clean_text(
            _pick_first(payload, "name")
            or _pick_first(company_full, "name", "brandName")
            or _pick_first(brand, "brandName")
        ) or "",
        "company_type": _clean_text(_pick_first(company_full, "typeName") or _pick_first(brand, "brandIndustry")),
        "industry": _clean_text(_pick_first(brand, "industryName") or _pick_first(company_full, "industry")),
        "company_size": _clean_text(_pick_first(brand, "scaleName") or _pick_first(company_full, "scaleName")),
        "financing_stage": _clean_text(_pick_first(brand, "stageName") or _pick_first(company_full, "stageName")),
        "city": _clean_text(_pick_first(company_full, "cityName", "city")),
        "address": _clean_text(_pick_first(company_full, "address", "addressInfo")),
        "website": _clean_text(_pick_first(company_full, "website")),
        "logo_url": _clean_text(_pick_first(company_full, "logo", "brandLogo") or _pick_first(brand, "logo", "brandLogo")),
        "description": _clean_text(
            _pick_first(company_full, "introduce", "introduction", "companyDesc")
            or _pick_first(brand, "introduce")
        ),
    }


def _extract_qcwy_fields(raw: dict[str, Any], company_id: str) -> dict[str, Any]:
    financing = raw.get("financingStage") or {}
    coinfo = raw.get("coinfo") if isinstance(raw.get("coinfo"), dict) else {}

    return {
        "source_company_id": normalize_company_id(
            "qcwy",
            company_id or _pick_first(raw, "companyId", "coId") or _nested_get(raw, "coinfo", "coid"),
        ),
        "company_name": _clean_text(
            _pick_first(raw, "companyName", "fullCompanyName", "companyNameEn")
            or _pick_first(coinfo, "coname", "brandName")
        ) or "",
        "company_type": _clean_text(_pick_first(raw, "companyTypeString", "orgTypeName") or _pick_first(coinfo, "cotype")),
        "industry": _clean_text(
            _pick_first(raw, "industryName", "companyIndustryType1Str")
            or _pick_first(coinfo, "indtype1", "indtype2", "coIndustryText")
        ),
        "company_size": _clean_text(
            _pick_first(raw, "companySizeString", "companySize", "orgSizeName")
            or _pick_first(coinfo, "cosize")
        ),
        "financing_stage": _clean_text(_pick_first(financing, "name") or _pick_first(raw, "financingStageName")),
        "city": _clean_text(_pick_first(raw, "cityName", "jobAreaString", "workCity") or _pick_first(coinfo, "areaString")),
        "address": _clean_text(
            _pick_first(raw, "address", "location")
            or _nested_get(raw, "workLocation", "workAddress")
            or _pick_first(coinfo, "caddr")
        ),
        "website": _clean_text(_pick_first(raw, "companyUrl", "companyHref") or _pick_first(coinfo, "webUrl")),
        "logo_url": _clean_text(_pick_first(raw, "companyLogo") or _pick_first(coinfo, "logourl")),
        "description": _clean_text(
            _pick_first(raw, "companyDesc", "company_desc", "description")
            or _nested_get(raw, "campusRootOrgInfo", "description")
            or _pick_first(coinfo, "coinfo")
        ),
    }


def _extract_zhilian_fields(raw: dict[str, Any], company_id: str) -> dict[str, Any]:
    data = raw.get("data") if isinstance(raw.get("data"), dict) else raw
    company_base = data.get("companyBase") or {}
    detailed_company = data.get("detailedCompany") or {}

    return {
        "source_company_id": normalize_company_id(
            "zhilian",
            company_id
            or _pick_first(company_base, "companyNumber", "number")
            or _pick_first(detailed_company, "companyNumber", "number"),
        ),
        "company_name": _clean_text(_pick_first(company_base, "companyName") or _pick_first(data, "companyName")) or "",
        "company_type": _clean_text(
            _pick_first(company_base, "companyTypeName", "companyType")
            or _pick_first(detailed_company, "companyTypeName")
        ),
        "industry": _clean_text(_pick_first(company_base, "industryName") or _pick_first(detailed_company, "industryName")),
        "company_size": _clean_text(
            _pick_first(company_base, "companySize", "companySizeString")
            or _pick_first(detailed_company, "companySize")
        ),
        "financing_stage": _clean_text(
            _pick_first(company_base, "financingStageName")
            or _nested_get(company_base, "financingStage", "name")
            or _nested_get(detailed_company, "financingStage", "name")
        ),
        "city": _clean_text(_pick_first(company_base, "cityName") or _pick_first(detailed_company, "cityName")),
        "address": _clean_text(_pick_first(company_base, "address") or _pick_first(detailed_company, "address")),
        "website": _clean_text(_pick_first(company_base, "companyUrl", "website")),
        "logo_url": _clean_text(_pick_first(company_base, "logoUrl", "companyLogo")),
        "description": _clean_text(
            _pick_first(company_base, "companyDescWithHtml", "companyDesc")
            or _pick_first(detailed_company, "companyDescription", "companyDesc")
        ),
    }


def extract_company_fields(source: str, raw: dict[str, Any], company_id: str) -> dict[str, Any]:
    if source == "boss":
        return _extract_boss_fields(raw, company_id)
    if source == "qcwy":
        return _extract_qcwy_fields(raw, company_id)
    if source == "zhilian":
        return _extract_zhilian_fields(raw, company_id)
    raise ValueError(f"unsupported source: {source}")


class CompanyStorageService:
    @staticmethod
    def company_model(source: str) -> Type[BaseCompanyModel]:
        return _model_for_source(source)

    async def get_existing_company_ids(self, source: str, company_ids: Iterable[str]) -> set[str]:
        normalized_ids = [normalize_company_id(source, item) for item in company_ids if item]
        if not normalized_ids:
            return set()
        model = self.company_model(source)
        rows = await model.filter(source_company_id__in=normalized_ids).values_list("source_company_id", flat=True)
        return set(rows)

    async def get_all_company_ids(self, source: str) -> set[str]:
        """获取该平台所有已入库的公司 ID（用于 ClickHouse 查询排除）"""
        model = self.company_model(source)
        rows = await model.all().values_list("source_company_id", flat=True)
        return set(rows)

    async def get_existing_queue_ids(self, source: str, company_ids: Iterable[str]) -> set[str]:
        normalized_ids = [normalize_company_id(source, item) for item in company_ids if item]
        if not normalized_ids:
            return set()
        rows = await CompanyCleaningQueue.filter(source=source, company_id__in=normalized_ids).values_list("company_id", flat=True)
        return set(rows)

    async def enqueue_company(self, source: str, company_id: str, company_name: str = "") -> tuple[CompanyCleaningQueue, bool]:
        normalized_id = normalize_company_id(source, company_id)
        defaults = {
            "company_name": company_name or "",
            "status": "pending",
            "error_msg": "",
            "retry_count": 0,
            "started_at": None,
            "finished_at": None,
            "jobs_fetched": 0,
            "jobs_stored": 0,
            "jobs_duplicate": 0,
            "jobs_failed": 0,
            "jobs_error_msg": "",
        }
        queue, created = await CompanyCleaningQueue.get_or_create(
            source=source,
            company_id=normalized_id,
            defaults=defaults,
        )
        if not created and company_name and queue.company_name != company_name:
            queue.company_name = company_name
            await queue.save(update_fields=["company_name", "updated_at"])
        return queue, created

    async def enqueue_companies(self, source: str, companies: Iterable[dict[str, str]]) -> int:
        created_count = 0
        for item in companies:
            _, created = await self.enqueue_company(
                source=source,
                company_id=item.get("company_id", ""),
                company_name=item.get("company_name", "") or "",
            )
            if created:
                created_count += 1
        return created_count

    async def get_company_record(self, source: str, company_id: str) -> Optional[BaseCompanyModel]:
        normalized_id = normalize_company_id(source, company_id)
        model = self.company_model(source)
        return await model.get_or_none(source_company_id=normalized_id)

    async def upsert_company(
        self,
        source: str,
        raw_data: dict[str, Any],
        *,
        company_id: Optional[str] = None,
    ) -> dict[str, Any]:
        normalized_id = normalize_company_id(source, company_id or "")
        fields = extract_company_fields(source, raw_data, normalized_id)
        normalized_id = fields["source_company_id"]
        if not normalized_id:
            raise ValueError(f"missing normalized company id for source={source}")
        if not fields["company_name"]:
            raise ValueError(f"missing company name for source={source} company_id={normalized_id}")

        model = self.company_model(source)
        record = await model.get_or_none(source_company_id=normalized_id)
        now = datetime.now()
        payload = {
            **fields,
            "raw_json": raw_data,
            "last_crawled_at": now,
        }

        if record:
            for key, value in payload.items():
                setattr(record, key, value)
            await record.save()
            created = False
        else:
            record = await model.create(
                **payload,
                first_crawled_at=now,
            )
            created = True

        return {
            "success": True,
            "created": created,
            "company_id": normalized_id,
            "company_name": record.company_name,
            "data_summary": {
                "source": source,
                "company_id": normalized_id,
                "company_name": record.company_name,
                "created": created,
            },
            "record": record,
        }

    async def mark_queue_processing(self, queue: CompanyCleaningQueue) -> None:
        queue.status = "processing"
        queue.error_msg = ""
        queue.started_at = datetime.now()
        queue.finished_at = None
        queue.jobs_fetched = 0
        queue.jobs_stored = 0
        queue.jobs_duplicate = 0
        queue.jobs_failed = 0
        queue.jobs_error_msg = ""
        await queue.save(
            update_fields=[
                "status",
                "error_msg",
                "started_at",
                "finished_at",
                "jobs_fetched",
                "jobs_stored",
                "jobs_duplicate",
                "jobs_failed",
                "jobs_error_msg",
                "updated_at",
            ]
        )

    async def mark_queue_result(
        self,
        queue: CompanyCleaningQueue,
        *,
        status: str,
        error_msg: str = "",
        increment_retry: bool = False,
        jobs_summary: Optional[dict[str, Any]] = None,
    ) -> None:
        queue.status = status
        queue.error_msg = error_msg or ""
        queue.finished_at = datetime.now()
        if jobs_summary:
            queue.jobs_fetched = int(jobs_summary.get("jobs_fetched") or 0)
            queue.jobs_stored = int(jobs_summary.get("stored_success") or 0)
            queue.jobs_duplicate = int(jobs_summary.get("duplicate") or 0)
            queue.jobs_failed = int(jobs_summary.get("failed") or 0)
            queue.jobs_error_msg = jobs_summary.get("error") or ""
        if increment_retry:
            queue.retry_count += 1
        await queue.save(
            update_fields=[
                "company_name",
                "status",
                "error_msg",
                "retry_count",
                "finished_at",
                "jobs_fetched",
                "jobs_stored",
                "jobs_duplicate",
                "jobs_failed",
                "jobs_error_msg",
                "updated_at",
            ]
        )


company_storage = CompanyStorageService()
