"""Shared WebHDFS helpers — resolve datanode redirects via NameNode JMX."""

from __future__ import annotations

import os
import time
from typing import Any
from urllib.parse import quote, urlparse, urlunparse

import httpx

HDFS_NN_URL = os.getenv("HDFS_NN_URL", "http://10.0.21.61:9870").rstrip("/")
HDFS_USER = os.getenv("HDFS_USER", "hdfs")
WEBHDFS = f"{HDFS_NN_URL}/webhdfs/v1"

_dn_cache: dict[str, Any] = {"ts": 0.0, "map": {}}


def encode_path(path: str) -> str:
    p = path if path.startswith("/") else f"/{path}"
    return "/".join(quote(seg, safe="") for seg in p.split("/"))


def datanode_host_map() -> dict[str, str]:
    now = time.time()
    if _dn_cache["map"] and now - _dn_cache["ts"] < 120:
        return _dn_cache["map"]
    mapping: dict[str, str] = {}
    try:
        with httpx.Client(timeout=8.0) as client:
            r = client.get(f"{HDFS_NN_URL}/jmx?qry=Hadoop:service=NameNode,name=NameNodeInfo")
            raw = (r.json().get("beans") or [{}])[0].get("LiveNodes") or "{}"
            if isinstance(raw, str):
                import json
                nodes = json.loads(raw)
            else:
                nodes = raw
            for key, info in nodes.items():
                host = key.split(":")[0]
                info_addr = (info or {}).get("infoAddr") or ""
                if info_addr:
                    mapping[host] = info_addr  # ip:9864
    except Exception:
        pass
    _dn_cache["map"] = mapping
    _dn_cache["ts"] = now
    return mapping


def resolve_redirect(url: str) -> str:
    parsed = urlparse(url)
    host = parsed.hostname or ""
    if not host or host.replace(".", "").isdigit():
        return url
    dn = datanode_host_map().get(host)
    if not dn:
        return url
    ip, _, port = dn.partition(":")
    port = parsed.port or port or "9864"
    return urlunparse(parsed._replace(netloc=f"{ip}:{port}"))


def open_bytes(path: str, max_bytes: int = 8_000_000) -> bytes:
    url = f"{WEBHDFS}{encode_path(path)}"
    with httpx.Client(timeout=90.0, follow_redirects=False) as client:
        r = client.get(url, params={"op": "OPEN", "user.name": HDFS_USER})
        if r.status_code in (301, 302, 307, 308):
            loc = r.headers.get("location") or ""
            if loc:
                r = client.get(resolve_redirect(loc), follow_redirects=True)
        elif r.status_code == 200 and r.headers.get("content-type", "").startswith("application/json"):
            loc = r.json().get("Location") or r.json().get("location") or ""
            if loc:
                r = client.get(resolve_redirect(loc), follow_redirects=True)
        if r.status_code >= 400:
            raise RuntimeError(r.text[:300])
        return r.content[:max_bytes]