"""LLM-based reranker for B2B account matching.

Uses Vertex AI with Gemini 3 Flash exclusively. Ollama support has been removed.
"""
from __future__ import annotations

import asyncio
import json
import logging
import os
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


def _mask_email_user(v: str) -> str:
    """Return masked email (***@domain) or the original string if not an email."""
    try:
        s = (v or "").strip()
        if "@" not in s:
            return s
        _user, domain = s.split("@", 1)
        return "***@" + domain
    except Exception:
        return v or ""


def _build_prompt(source: Dict[str, Any], candidates: List[Dict[str, Any]]) -> str:
    def _sv(d: Dict[str, Any], k: str) -> str:
        v = d.get(k)
        return "" if v is None else str(v)

    # Mask emails in source for privacy
    src_email = _mask_email_user(_sv(source, "Email"))

    src_txt = (
        f"Source:\n"
        f"- Company: {_sv(source,'Company')}\n"
        f"- Domain: {_sv(source,'Domain')}\n"
        f"- Email: {src_email}\n"
        f"- Name: {_sv(source,'Name')}\n\n"
    )
    cands = []
    for c in candidates:
        # Derive fields
        cid = _sv(c, "id")
        cname = (
            _sv(c, "name")
            or _sv(c, "Company")
            or _sv(c, "Account Name")
            or _sv(c, "Name")
        )
        cdomain = (
            _sv(c, "domain")
            or _sv(c, "Domain")
            or _sv(c, "Website")
            or _sv(c, "EmailDomain")
        )
        calt = c.get("alt") or []
        cupa = _sv(c, "upa")

        # Optional: mask any email-like values in candidate fields
        if "@" in cdomain:
            cdomain = _mask_email_user(cdomain)
        if isinstance(calt, list):
            calt = [
                _mask_email_user(x) if isinstance(x, str) and ("@" in x) else x
                for x in calt
            ]
        if "@" in cupa:
            cupa = _mask_email_user(cupa)

        cands.append(
            {
                "id": cid,
                "name": cname,
                "domain": cdomain,
                "alt": calt,
                "upa": cupa,
            }
        )
    cand_txt = json.dumps(cands, ensure_ascii=False)

    return (
        "You are a strict re-ranker for B2B account matching.\n"
        "Choose the best candidate by business identity (domain + name), or NONE.\n"
        "Use ONLY the provided fields. Do not invent or copy-edit data.\n"
        "Respond as compact JSON, no prose.\n\n"
        + src_txt
        + "Candidates (pick one by id or NONE):\n"
        + cand_txt
        + "\n\n"
        'Return JSON: {"choice_id":"<id|NONE>","reason":"<max 12 words>"}'
    )


async def choose(
    source: Dict[str, Any], candidates: List[Dict[str, Any]], *, timeout: float = 2.5
) -> Optional[Dict[str, str]]:
    """Call Vertex AI to pick best candidate or NONE.

    Returns: {"choice_id": str, "reason": str} | None
    """
    # Check if reranking is enabled
    if os.getenv("ENABLE_OLLAMA_RERANK", "0") != "1" and os.getenv("ENABLE_LLM_RERANK", "0") != "1":
        return None

    prompt = _build_prompt(source, candidates)

    # Try Vertex AI first (preferred)
    try:
        import vertexai
        from vertexai.generative_models import GenerativeModel, GenerationConfig

        project = (
            os.getenv("VERTEX_PROJECT")
            or os.getenv("GOOGLE_CLOUD_PROJECT")
            or os.getenv("GCP_PROJECT")
        )
        location = os.getenv("VERTEX_LOCATION", "us-central1")
        model_name = os.getenv("VERTEX_MODEL") or os.getenv("RERANK_MODEL") or "gemini-3-flash-preview"

        # Route Gemini 3 to global
        if model_name.lower().startswith("gemini-3-"):
            location = "global"

        if project:
            vertexai.init(project=project, location=location)
            model = GenerativeModel(model_name)
            config = GenerationConfig(
                temperature=0.1,
                max_output_tokens=128,
                response_mime_type="application/json",
            )

            def _call():
                return model.generate_content([prompt], generation_config=config)

            response = await asyncio.wait_for(
                asyncio.to_thread(_call),
                timeout=timeout
            )

            txt = getattr(response, "text", None) or "{}"
            try:
                out = json.loads(txt)
            except Exception:
                logger.debug("Rerank JSON parse failed: %s", txt[:100])
                return None

            cid = out.get("choice_id")
            reason = out.get("reason")
            if not isinstance(cid, str) or not isinstance(reason, str):
                return None
            return {"choice_id": cid.strip(), "reason": reason.strip()[:80]}

    except ImportError:
        logger.debug("Vertex AI SDK not available for rerank")
    except asyncio.TimeoutError:
        logger.debug("Rerank timed out")
        return None
    except Exception as e:
        logger.debug("Vertex AI rerank failed: %s", e)

    # Ollama support has been removed - Vertex AI is required
    if os.getenv("LLM_PROVIDER", "").lower() == "ollama":
        logger.error(
            "Ollama reranking has been removed. Configure Vertex AI instead: "
            "set VERTEX_PROJECT and optionally VERTEX_MODEL (default: gemini-3-flash-preview)"
        )

    return None
