backend/app/clients/hana_client.py

from __future__ import annotations

import json
import logging
import time
from typing import Any

import httpx

from app.config import settings

LOG = logging.getLogger(__name__)


class HanaClient:
    """Handles HANA API authentication, model/persona discovery, and inference."""

    def __init__(self) -> None:
        self._base = settings.hana_base_url.rstrip("/")
        self._access_token: str = ""
        self._refresh_token: str = ""
        self._token_expiry: float = 0
        self._klatchat_access_token: str = ""
        self._klatchat_refresh_token: str = ""
        self._klatchat_token_expiry: float = 0
        self._client = httpx.AsyncClient(timeout=30.0)
        self._persona_cache: dict[str, str | None] = {}

    def _uses_klatchat_model(self, model_id: str) -> bool:
        """Legacy HANA-JWT path for Security — disabled when using direct vLLM (brainforge-webapp pattern)."""
        if settings._neon_security_direct_vllm_enabled(model_id):
            return False
        if not (settings.hana_password_klatchat or "").strip():
            return False
        mid = (model_id or "").lower()
        return "brainforge/security" in mid or "/security@" in mid

    async def authenticate(self) -> None:
        resp = await self._client.post(
            f"{self._base}/auth/login",
            json={
                "username": settings.hana_username,
                "password": settings.hana_password,
                "token_name": "LLMChats3",
                "client_id": "llm-chat-tool",
            },
        )
        resp.raise_for_status()
        data = resp.json()
        self._access_token = data["access_token"]
        self._refresh_token = data["refresh_token"]
        self._token_expiry = data.get("expiration", time.time() + 3600)
        LOG.info("HANA auth success (user=%s)", data.get("username"))

    async def authenticate_klatchat(self) -> None:
        """Login for BrainForge/Security (4090) using HANA_KLATCHAT_PASSWORD."""
        uname = (settings.hana_username_klatchat or settings.hana_username).strip()
        pwd = (settings.hana_password_klatchat or "").strip()
        if not pwd:
            raise ValueError("hana_password_klatchat not set")
        resp = await self._client.post(
            f"{self._base}/auth/login",
            json={
                "username": uname,
                "password": pwd,
                "token_name": "LLMChats3",
                "client_id": "llm-chat-tool",
            },
        )
        resp.raise_for_status()
        data = resp.json()
        self._klatchat_access_token = data["access_token"]
        self._klatchat_refresh_token = data["refresh_token"]
        self._klatchat_token_expiry = data.get("expiration", time.time() + 3600)
        LOG.info("HANA klatchat auth success (user=%s)", data.get("username"))

    async def _ensure_token(self) -> None:
        if time.time() >= self._token_expiry - 60:
            try:
                resp = await self._client.post(
                    f"{self._base}/auth/refresh",
                    json={
                        "access_token": self._access_token,
                        "refresh_token": self._refresh_token,
                    },
                )
                resp.raise_for_status()
                data = resp.json()
                self._access_token = data["access_token"]
                self._refresh_token = data["refresh_token"]
                self._token_expiry = data.get("expiration", time.time() + 3600)
                LOG.info("HANA token refreshed")
            except Exception:
                LOG.warning("Token refresh failed, re-authenticating")
                await self.authenticate()

    async def _ensure_klatchat_token(self) -> None:
        if not (settings.hana_password_klatchat or "").strip():
            await self._ensure_token()
            return
        if not self._klatchat_access_token or time.time() >= self._klatchat_token_expiry - 60:
            try:
                resp = await self._client.post(
                    f"{self._base}/auth/refresh",
                    json={
                        "access_token": self._klatchat_access_token,
                        "refresh_token": self._klatchat_refresh_token,
                    },
                )
                resp.raise_for_status()
                data = resp.json()
                self._klatchat_access_token = data["access_token"]
                self._klatchat_refresh_token = data["refresh_token"]
                self._klatchat_token_expiry = data.get("expiration", time.time() + 3600)
                LOG.info("HANA klatchat token refreshed")
            except Exception:
                LOG.warning("Klatchat token refresh failed, re-authenticating")
                await self.authenticate_klatchat()

    @property
    def _headers(self) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {self._access_token}",
            "Content-Type": "application/json",
        }

    def _headers_klatchat(self) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {self._klatchat_access_token}",
            "Content-Type": "application/json",
        }

    def _headers_for_model(self, model_id: str) -> dict[str, str]:
        if self._uses_klatchat_model(model_id) and (settings.hana_password_klatchat or "").strip():
            return self._headers_klatchat()
        return self._headers

    async def _ensure_headers_for_model(self, model_id: str) -> dict[str, str]:
        if self._uses_klatchat_model(model_id) and (settings.hana_password_klatchat or "").strip():
            await self._ensure_klatchat_token()
            return self._headers_klatchat()
        await self._ensure_token()
        return self._headers

    def _parse_models_payload(self, data: dict[str, Any]) -> list[dict[str, Any]]:
        models: list[dict[str, Any]] = []
        for m in data.get("models", []):
            model_id = f"{m['name']}@{m['version']}"
            personas = []
            for p in m.get("personas", []):
                pname = p.get("persona_name", "")
                sp = p.get("system_prompt") or ""
                cache_key = f"{model_id}:{pname}"
                self._persona_cache[cache_key] = sp if sp else None
                personas.append({
                    "id": p.get("id", p.get("persona_name", "")),
                    "persona_name": pname,
                    "description": p.get("description"),
                    "system_prompt": p.get("system_prompt"),
                    "enabled": p.get("enabled", True),
                })
            models.append({
                "name": m["name"],
                "version": m["version"],
                "model_id": model_id,
                "personas": personas,
            })
        return models

    async def _merge_direct_vllm_security_models(
        self, models: list[dict[str, Any]], seen: set[str]
    ) -> None:
        """Merge models from OpenAI-compatible GET /v1/models (NeonGeckoCom/brainforge-webapp pattern)."""
        base = (settings.neon_security_vllm_base_url or "").strip().rstrip("/")
        key = (settings.vllm_api_key or "").strip()
        if not base or not key:
            return
        url = f"{base}/v1/models"
        try:
            resp = await self._client.get(
                url,
                headers={"Authorization": f"Bearer {key}"},
            )
            if resp.status_code != 200:
                LOG.warning("Neon direct vLLM %s -> %s", url, resp.status_code)
                return
            payload = resp.json()
            added = 0
            for o in payload.get("data", []) or []:
                mid = (o.get("id") or "").strip()
                if not mid or "@" not in mid:
                    continue
                if mid in seen:
                    continue
                name, version = mid.split("@", 1)
                vanilla = {
                    "id": "vanilla",
                    "persona_name": "vanilla",
                    "description": None,
                    "system_prompt": None,
                    "enabled": True,
                }
                self._persona_cache[f"{mid}:vanilla"] = None
                models.append({
                    "name": name,
                    "version": version,
                    "model_id": mid,
                    "personas": [vanilla],
                })
                seen.add(mid)
                added += 1
            if added:
                LOG.info("Merged %s model(s) from direct vLLM %s", added, base)
            # region agent log
            try:
                with open(r"c:\Users\dream\CCAI-Demo-FEAT_Config\debug-c86901.log", "a", encoding="utf-8") as _df:
                    _df.write(
                        json.dumps(
                            {
                                "sessionId": "c86901",
                                "hypothesisId": "H5",
                                "location": "LLMChats3/hana_client._merge_direct_vllm",
                                "message": "direct_vllm_models",
                                "data": {
                                    "app": "LLMChats3",
                                    "url": url,
                                    "http_status": resp.status_code,
                                    "added_count": added,
                                },
                                "timestamp": int(time.time() * 1000),
                            }
                        )
                        + "\n"
                    )
            except Exception:
                pass
            # endregion
        except Exception as exc:
            LOG.warning("Neon direct vLLM model merge failed: %s", exc)

    async def _enrich_security_personas_from_hana(self, models: list[dict[str, Any]]) -> None:
        """Replace Security personas with HANA get_personas when the server allows (persona prompts live on HANA)."""
        for m in models:
            mid = m.get("model_id") or ""
            if "security" not in mid.lower():
                continue
            try:
                plist = await self.get_personas(mid)
                if plist:
                    m["personas"] = plist
                    for p in plist:
                        pname = p.get("persona_name", "")
                        sp = p.get("system_prompt") or ""
                        self._persona_cache[f"{mid}:{pname}"] = sp if sp else None
                    LOG.info("Enriched Security personas from HANA get_personas for %s", mid)
            except Exception as exc:
                LOG.debug("HANA get_personas enrichment skipped for %s: %s", mid, exc)

    async def _ensure_security_model_stub_if_missing(self, models: list[dict[str, Any]]) -> None:
        """If vLLM merge failed (401) and HANA supplement failed, still list Security so the UI can show it."""
        sid = "BrainForge/Security@2026.03.18"
        wanted = [x.strip() for x in (settings.hana_neon_model_supplement_ids or "").split(",") if x.strip()]
        if not any((w == sid or ("security" in w.lower() and "@" in w)) for w in wanted):
            return
        if any((m.get("model_id") == sid) for m in models):
            return
        try:
            plist = await self.get_personas(sid)
            name, ver = sid.split("@", 1)
            for p in plist or []:
                pname = p.get("persona_name", "")
                sp = p.get("system_prompt") or ""
                self._persona_cache[f"{sid}:{pname}"] = sp if sp else None
            models.append({
                "name": name,
                "version": ver,
                "model_id": sid,
                "personas": plist or [],
            })
            LOG.info("Security model added via HANA get_personas (stub path)")
        except Exception as exc:
            LOG.warning(
                "Security model not in HANA/vLLM responses (%s); adding minimal stub entry.",
                exc,
            )
            vanilla = {
                "id": "vanilla",
                "persona_name": "vanilla",
                "description": "Set VLLM_API_KEY to the 4090-x1-3 vLLM Bearer token (same as HF Space API_KEY) so vLLM discovery works; persona prompts come from HANA when allowed.",
                "system_prompt": None,
                "enabled": True,
            }
            self._persona_cache[f"{sid}:vanilla"] = None
            models.append({
                "name": "BrainForge/Security",
                "version": "2026.03.18",
                "model_id": sid,
                "personas": [vanilla],
            })

    async def get_models(self) -> list[dict[str, Any]]:
        await self._ensure_token()
        resp = await self._client.post(
            f"{self._base}/brainforge/get_models",
            headers=self._headers,
            json={},
        )
        resp.raise_for_status()
        models = self._parse_models_payload(resp.json())
        seen = {m["model_id"] for m in models}

        # The direct-vLLM merge is the HuggingFace-deployment fallback path:
        # it talks directly to 4090-x1-3 because the HF runtime can't reach
        # HANA. In local dev, HANA itself returns BrainForge/Security in the
        # standard get_models response and proxies inference to 4090-x1-3
        # for us, so the merge is redundant. Skip it whenever HANA already
        # gave us Security. (Gated on VLLM_API_KEY, the Bearer token for
        # the direct vLLM endpoint - distinct from the HANA login.)
        hana_has_security = any(
            "security" in (mid or "").lower() for mid in seen
        )
        if (
            (settings.vllm_api_key or "").strip()
            and (settings.neon_security_vllm_base_url or "").strip()
            and not hana_has_security
        ):
            await self._merge_direct_vllm_security_models(models, seen)

        await self._append_supplement_models(models)
        await self._enrich_security_personas_from_hana(models)
        await self._ensure_security_model_stub_if_missing(models)
        # region agent log
        try:
            mids = [m.get("model_id", "") for m in models]
            with open(r"c:\Users\dream\CCAI-Demo-FEAT_Config\debug-c86901.log", "a", encoding="utf-8") as _df:
                _df.write(
                    json.dumps(
                        {
                            "sessionId": "c86901",
                            "hypothesisId": "H3",
                            "location": "LLMChats3/hana_client.get_models",
                            "message": "merged_models",
                            "data": {
                                "app": "LLMChats3",
                                "model_count": len(models),
                                "has_security_in_list": any("security" in (x or "").lower() for x in mids),
                                "model_ids_tail": mids[-8:],
                            },
                            "timestamp": int(time.time() * 1000),
                        }
                    )
                    + "\n"
                )
        except Exception:
            pass
        # endregion
        return models

    async def get_personas(self, model_id: str) -> list[dict[str, Any]]:
        """Fetch personas for a model_id (used when get_models omits a model)."""
        hdrs = await self._ensure_headers_for_model(model_id)
        resp = await self._client.post(
            f"{self._base}/brainforge/get_personas",
            headers=hdrs,
            json={"model_id": model_id},
        )
        resp.raise_for_status()
        data = resp.json()
        out = []
        for p in data.get("personas", []):
            pname = p.get("persona_name", "")
            sp = p.get("system_prompt") or ""
            cache_key = f"{model_id}:{pname}"
            self._persona_cache[cache_key] = sp if sp else None
            out.append({
                "id": p.get("id", p.get("persona_name", "")),
                "persona_name": pname,
                "description": p.get("description"),
                "system_prompt": p.get("system_prompt"),
                "enabled": p.get("enabled", True),
            })
        return out

    async def _append_supplement_models(self, models: list[dict[str, Any]]) -> None:
        """Merge models listed in settings but missing from get_models (HANA may omit some)."""
        raw = (settings.hana_neon_model_supplement_ids or "").strip()
        extras = [x.strip() for x in raw.split(",") if x.strip()]
        seen = {m["model_id"] for m in models}
        for mid in extras:
            if mid in seen:
                continue
            if "@" not in mid:
                LOG.warning("Invalid supplement model_id (expected name@version): %s", mid)
                continue
            try:
                personas = await self.get_personas(mid)
                name, version = mid.split("@", 1)
                models.append({
                    "name": name,
                    "version": version,
                    "model_id": mid,
                    "personas": personas,
                })
                seen.add(mid)
                LOG.info("Merged supplement Neon model from get_personas: %s", mid)
            except Exception as exc:
                LOG.debug("Supplement Neon model %s not merged: %s", mid, exc)

    def get_persona_system_prompt(self, model_id: str, persona_name: str) -> str | None:
        """Look up a persona's built-in system_prompt from the cache."""
        return self._persona_cache.get(f"{model_id}:{persona_name}")

    async def get_inference(
        self,
        query: str,
        model_id: str,
        persona_name: str,
        system_prompt: str | None = None,
        history: list[tuple[str, str]] | None = None,
        temperature: float = 0.7,
        max_tokens: int = 1024,
    ) -> dict[str, Any]:
        hdrs = await self._ensure_headers_for_model(model_id)
        hist = [[role, content] for role, content in (history or [])]

        persona_payload: dict[str, Any] = {"persona_name": persona_name}
        if system_prompt:
            persona_payload["system_prompt"] = system_prompt

        body = {
            "query": query,
            "history": hist,
            "persona": persona_payload,
            "model": model_id,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "extra_body": {},
            "llm_name": model_id.split("@")[0] if "@" in model_id else model_id,
            "llm_revision": model_id.split("@")[1] if "@" in model_id else "",
        }

        t0 = time.time()
        resp = await self._client.post(
            f"{self._base}/brainforge/get_inference",
            headers=hdrs,
            json=body,
        )
        elapsed = time.time() - t0
        resp.raise_for_status()
        data = resp.json()
        return {
            "response": data.get("response", ""),
            "elapsed_seconds": round(elapsed, 2),
            "finish_reason": data.get("finish_reason", "stop"),
        }

    async def close(self) -> None:
        await self._client.aclose()


hana_client = HanaClient()