"""OCR 面单识别响应模型 + 国家代码映射。"""

from __future__ import annotations

from typing import Any

from pydantic import BaseModel, ConfigDict


class OcrResult(BaseModel):
    """单张面单 OCR 识别结果。"""

    model_config = ConfigDict(extra="allow")

    file: str | None = None
    layout: str | None = None
    courier: str | None = None
    courier_name: str | None = None
    country: str | None = None
    error: str | None = None
    tracking: str | None = None
    s_name: str | None = None
    s_street: str | None = None
    s_city: str | None = None
    s_state: str | None = None
    s_zip: str | None = None
    r_name: str | None = None
    r_street: str | None = None
    r_city: str | None = None
    r_state: str | None = None
    r_zip: str | None = None
    marker: str | None = None

    @property
    def ok(self) -> bool:
        return not self.error


# OCR 服务已有分类器模型的国家(版面见 api.md §4)
SUPPORTED_COUNTRIES = {"US", "UK", "DE", "MX", "SA", "CA", "ES", "IT"}

# ISO 国家代码 / 中文国名 → OCR 国家代码
_COUNTRY_MAP: dict[str, str] = {
    # 代码
    "GB": "UK", "UK": "UK", "US": "US", "DE": "DE", "MX": "MX", "SA": "SA",
    "CA": "CA", "ES": "ES", "IT": "IT",
    # 中文名
    "英国": "UK", "美国": "US", "德国": "DE", "墨西哥": "MX", "沙特": "SA",
    "加拿大": "CA", "西班牙": "ES", "意大利": "IT",
}


def country_to_ocr_code(country: Any) -> str | None:
    """将 ISO 代码或中文国名转为 OCR 服务支持的国家代码,不支持则返回 None。"""
    if not country or not isinstance(country, str):
        return None
    code = _COUNTRY_MAP.get(country.strip().upper()) or _COUNTRY_MAP.get(country.strip())
    return code if code in SUPPORTED_COUNTRIES else None