#!/usr/bin/env python3
# #region check_llm_certs [C:4] [TYPE Module] [SEMANTICS ssl, certs, diagnostics, python]
# @BRIEF SSL diagnostics in Python — tests all real libraries used by ss-tools.
#   Reads TARGET_URL from env or CLI --target.
#   Uses capath by default (cafile is unreliable — see ADR-0009).
# @USAGE ./scripts/check_llm_certs.py [--target https://...]
# @LAYER Infrastructure
# #endregion check_llm_certs

from __future__ import annotations

import argparse
import json
import os
import platform
import subprocess
import sys
import tempfile
import traceback
from pathlib import Path

# ── Configuration defaults ──
CA_BUNDLE = Path("/etc/ssl/certs/ca-certificates.crt")
CA_DIR = Path("/etc/ssl/certs")
CUSTOM_DIRS = [
    Path("/usr/local/share/ca-certificates/custom"),
    Path("/usr/local/share/ca-certificates/llm"),
]
NSS_DB_DIR = Path(os.environ.get("HOME", "/root")) / ".pki" / "nssdb"

results: dict = {}


# ── Lazy library loader ──
_LIBS: dict = {}

def _import(name: str) -> str | None:
    if name in _LIBS:
        return _LIBS[name] if isinstance(_LIBS[name], str) else None
    try:
        __import__(name)
        _LIBS[name] = True
        return None
    except ImportError:
        _LIBS[name] = "NOT_INSTALLED"
        return _LIBS[name]

def _lib_ver(name: str) -> str:
    err = _import(name)
    if err:
        return err
    mod = __import__(name)
    v = getattr(mod, "__version__", None)
    return str(v) if v else "installed"

def _count_certs(path: str) -> int:
    try:
        with open(path, "rb") as f:
            return f.read().count(b"-----BEGIN CERTIFICATE-----")
    except Exception:
        return 0


# ── Certificate info with proper in_bundle check ──

def _cert_info(path: Path) -> dict:
    """Extract cert info. Tries PEM first, then DER."""
    info = {"file": str(path), "format": path.suffix}

    def _try_read(fmt: str) -> dict | None:
        args = ["openssl", "x509", "-in", str(path), "-noout",
                "-subject", "-issuer", "-fingerprint", "-sha256"]
        if fmt == "DER":
            args += ["-inform", "DER"]
        r = subprocess.run(args, capture_output=True, text=True, timeout=5)
        if r.returncode != 0:
            return None
        res = {}
        for line in r.stdout.strip().split("\n"):
            if line.startswith("subject="): res["subject"] = line[8:]
            elif line.startswith("issuer="): res["issuer"] = line[7:]
            elif line.startswith("SHA256 Fingerprint="): res["fingerprint"] = line[19:]
        return res

    d = _try_read("PEM") or _try_read("DER") or {}
    info["subject"] = d.get("subject", "UNKNOWN_FORMAT")
    info["issuer"] = d.get("issuer", "?")
    info["fingerprint"] = d.get("fingerprint", "?")
    if _try_read("PEM") is None and d:
        info["format"] += " (DER)"

    # Proper in_bundle check: compare fingerprint against EACH cert in bundle
    fp = info["fingerprint"]
    info["in_bundle"] = "unknown"
    if fp and fp != "?" and CA_BUNDLE.exists():
        rc, out, _ = _run(
            ["awk", '-vcmd=openssl x509 -noout -fingerprint -sha256',
             '/BEGIN/{close(cmd)};{print | cmd}', str(CA_BUNDLE)], 10
        )
        if rc == 0:
            info["in_bundle"] = "yes" if fp in out else "no"
        else:
            # Fallback: simple file search
            with open(CA_BUNDLE, "rb") as f:
                info["in_bundle"] = "yes" if fp.encode() in f.read() else "no"
    return info


# ── OpenSSL s_client runner (with stdin EOF to prevent hang) ──

def _run_openssl(cmd: list[str], timeout: int = 20) -> tuple[int, str, str]:
    """Run openssl s_client with stdin=EOF to prevent hang after handshake."""
    full_cmd = cmd + ["-verify_return_error"]
    try:
        r = subprocess.run(full_cmd, input="", capture_output=True,
                           text=True, timeout=timeout)
        return r.returncode, r.stdout, r.stderr
    except FileNotFoundError:
        return -1, "", "command not found"
    except subprocess.TimeoutExpired:
        return -2, "", f"timeout ({timeout}s)"


def _run(cmd: list[str], timeout: int = 15) -> tuple[int, str, str]:
    """Generic subprocess runner."""
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
        return r.returncode, r.stdout, r.stderr
    except FileNotFoundError:
        return -1, "", "command not found"
    except subprocess.TimeoutExpired:
        return -2, "", f"timeout ({timeout}s)"


# ── OpenAI-compatible s_client parsing ──

def _parse_openssl(key: str, rc: int, out: str, err: str) -> None:
    for line in out.split("\n"):
        if "Verify return code" in line:
            results[key] = line.strip()
            emoji = "✅" if "0 (ok)" in line else "❌"
            print(f"  {key}: {emoji} {line.strip()}")
            return
    # If timed out, say so clearly
    if "timeout" in err:
        results[key] = f"TIMEOUT (proxy? check HTTP_PROXY env)"
        print(f"  {key}: ❌ {results[key]}")
        return
    msg = err.strip() or f"NO_OUTPUT (rc={rc})"
    results[key] = f"ERROR: {msg}"
    print(f"  {key}: ❌ {msg}")


# ── SSL context builders ──

def _ctx_cafile():
    import ssl
    return ssl.create_default_context(cafile=str(CA_BUNDLE))

def _ctx_capath():
    import ssl
    return ssl.create_default_context(capath=str(CA_DIR))

def _ctx_cafile_capath():
    import ssl
    ctx = ssl.create_default_context()
    if CA_BUNDLE.exists() and _count_certs(str(CA_BUNDLE)) > 0:
        ctx.load_verify_locations(cafile=str(CA_BUNDLE))
    if CA_DIR.is_dir():
        ctx.load_verify_locations(capath=str(CA_DIR))
    return ctx


# ── HTTP testers ──

def _test_httpx(key: str, verify) -> None:
    err = _import("httpx")
    if err:
        results[key] = err
        print(f"  {key}: ⚠ {err}")
        return
    import httpx
    try:
        r = httpx.get(TARGET_URL, verify=verify, timeout=10)
        results[key] = f"OK status={r.status_code}"
        print(f"  {key}: ✅ status={r.status_code}")
    except httpx.ConnectError as e:
        results[key] = f"FAIL ConnectError"
        print(f"  {key}: ❌ {e}")
    except Exception as e:
        results[key] = f"FAIL {type(e).__name__}"
        print(f"  {key}: ❌ {type(e).__name__}: {e}")

def _test_requests(key: str, verify) -> None:
    err = _import("requests")
    if err:
        results[key] = err
        print(f"  {key}: ⚠ {err}")
        return
    import requests as req
    try:
        r = req.get(TARGET_URL, verify=verify, timeout=10)
        results[key] = f"OK status={r.status_code}"
        print(f"  {key}: ✅ status={r.status_code}")
    except req.exceptions.SSLError as e:
        results[key] = f"FAIL SSLError"
        print(f"  {key}: ❌ {e}")
    except Exception as e:
        results[key] = f"FAIL {type(e).__name__}"
        print(f"  {key}: ❌ {type(e).__name__}: {e}")


def _section(name: str) -> None:
    print(f"\n=== {name} ===")


# ═══════════════════════════════════════════════════════
# DIAGNOSIS
# ═══════════════════════════════════════════════════════

def _diagnose(s: dict) -> str:
    issues = []
    for k in ("openssl_default", "openssl_cafile", "openssl_capath"):
        v = str(s.get(k, ""))
        issues.append(f"{k}={'OK' if '0 (ok)' in v else 'FAIL'}")
    for lib, methods in [("httpx", ("true", "cafile", "capath", "cafile_capath", "false")),
                          ("requests", ("true", "cafile", "capath", "false"))]:
        for method in methods:
            v = str(s.get(f"{lib}_{method}", ""))
            issues.append(f"{lib}_{method}={'OK' if v.startswith('OK') else 'FAIL'}")
    # Recommendation
    if s.get("httpx_cafile_capath", "").startswith("OK") and not s.get("httpx_cafile", "").startswith("OK"):
        issues.append("FIX=use_cafile_capath_for_httpx")
    elif s.get("httpx_capath", "").startswith("OK") and not s.get("httpx_cafile", "").startswith("OK"):
        issues.append("FIX=use_capath_for_httpx")
    elif s.get("requests_capath", "").startswith("OK") and not s.get("requests_cafile", "").startswith("OK"):
        issues.append("FIX=use_capath_for_requests")
    elif s.get("httpx_false", "").startswith("OK"):
        issues.append("FIX=llm_ssl_verify_false_works")
    else:
        issues.append("FIX=needs_deeper_investigation")
    return "; ".join(issues)


# ═══════════════════════════════════════════════════════
# MAIN
# ═══════════════════════════════════════════════════════

def main() -> None:
    global results, TARGET_URL
    results = {}

    parser = argparse.ArgumentParser(description="SSL Certificate Diagnostic Tool")
    parser.add_argument("--target", default=os.environ.get("TARGET_URL", "https://lite.ai.rusal.com"),
                        help="Target URL to test (default: https://lite.ai.rusal.com)")
    args = parser.parse_args()
    TARGET_URL = args.target
    host = TARGET_URL.replace("https://", "").split("/")[0]
    LLM_SSL_VERIFY = os.environ.get("LLM_SSL_VERIFY", "true")

    _section("1. System Information")
    results["target_url"] = TARGET_URL
    results["host"] = host
    results["llm_ssl_verify"] = LLM_SSL_VERIFY
    results["platform"] = platform.platform()
    results["python_version"] = sys.version
    results["httpx_version"] = _lib_ver("httpx")
    results["requests_version"] = _lib_ver("requests")
    results["openai_version"] = _lib_ver("openai")
    certifi_err = _import("certifi")
    if certifi_err:
        results["certifi_path"] = certifi_err
        results["certifi_certs"] = 0
    else:
        import certifi as _certifi
        certifi_path = _certifi.where()
        results["certifi_path"] = certifi_path
        results["certifi_certs"] = _count_certs(certifi_path)
    print(f"  target={TARGET_URL}  host={host}")
    print(f"  httpx={results['httpx_version']}  requests={results['requests_version']}  openai={results['openai_version']}")
    print(f"  certifi={results.get('certifi_path','?')} ({results.get('certifi_certs','?')} certs)")

    _section("2. System CA Store")
    results["ca_bundle"] = str(CA_BUNDLE)
    results["ca_bundle_certs"] = _count_certs(str(CA_BUNDLE)) if CA_BUNDLE.exists() else 0
    results["ca_hash_links"] = len(list(CA_DIR.glob("[0-9a-f]*.[0-9]"))) if CA_DIR.is_dir() else 0
    print(f"  bundle={CA_BUNDLE} ({results['ca_bundle_certs']} certs)")
    print(f"  capath={CA_DIR} ({results['ca_hash_links']} hash links)")

    print("  custom_certs:")
    custom_certs = []
    for d in CUSTOM_DIRS:
        if not d.is_dir():
            continue
        for f in sorted(d.iterdir()):
            if f.suffix in (".pem", ".crt") and f.is_file():
                ci = _cert_info(f)
                custom_certs.append(ci)
                print(f"    {ci['file']}")
                print(f"      subj={ci['subject']}  fmt={ci['format']}  in_bundle={ci['in_bundle']}")
    results["custom_certs"] = custom_certs
    results["custom_certs_count"] = len(custom_certs)

    _section("3. OpenSSL s_client")
    for label, extra_args in [
        ("openssl_default", []),
        ("openssl_cafile", ["-CAfile", str(CA_BUNDLE)]),
        ("openssl_capath", ["-CApath", str(CA_DIR)]),
    ]:
        cmd = ["openssl", "s_client", "-connect", f"{host}:443",
               "-servername", host] + extra_args
        rc, out, err = _run_openssl(cmd, 20)
        _parse_openssl(label, rc, out, err)

    # Chain test: combine PEM certs
    cert_files = []
    for d in CUSTOM_DIRS:
        if d.is_dir():
            for f in d.iterdir():
                if f.suffix in (".pem", ".crt") and f.is_file():
                    rc, _, _ = _run(["openssl", "x509", "-in", str(f), "-noout"], 3)
                    if rc == 0:
                        cert_files.append(f)
    if len(cert_files) >= 3:
        with tempfile.NamedTemporaryFile(mode="wb", suffix=".pem", delete=False) as tf:
            chain_path = tf.name
            for name in ["RUSAL_ROOT", "Policy_CA", "RGM_Issuing"]:
                for f in cert_files:
                    if name in f.name:
                        tf.write(f.read_bytes())
                        break
        rc, out, err = _run_openssl(
            ["openssl", "s_client", "-connect", f"{host}:443",
             "-CAfile", chain_path, "-servername", host], 20)
        os.unlink(chain_path)
        _parse_openssl("openssl_chain", rc, out, err)
    else:
        print("  openssl_chain: ⚠ < 3 PEM certs, skipping")

    _section("4. Python httpx (как LLMClient)")
    _test_httpx("httpx_true", verify=True)
    _test_httpx("httpx_cafile", verify=_ctx_cafile())
    _test_httpx("httpx_capath", verify=_ctx_capath())
    _test_httpx("httpx_cafile_capath", verify=_ctx_cafile_capath())
    _test_httpx("httpx_false", verify=False)

    _section("5. Python requests (как translate)")
    _test_requests("requests_true", verify=True)
    _test_requests("requests_cafile", verify=str(CA_BUNDLE))
    _test_requests("requests_capath", verify=str(CA_DIR))
    _test_requests("requests_false", verify=False)

    _section("6. NSS Chromium")
    rc, out, err = _run(["certutil", "-L", "-d", f"sql:{NSS_DB_DIR}"], 5)
    if rc == 0:
        results["nss_available"] = True
        results["nss_db"] = str(NSS_DB_DIR)
        lines = [l.strip() for l in out.split("\n")
                 if l.strip() and "Certificate Nickname" not in l]
        results["nss_certs_count"] = len(lines)
        print(f"  DB={NSS_DB_DIR}  certs={len(lines)}")
        for l in lines:
            if any(x in l.lower() for x in ("rusal", "llm-", "custom-")):
                print(f"    {l}")
    else:
        results["nss_available"] = False
        results["nss_note"] = err.strip() or "certutil or NSS DB unavailable"
        print(f"  NSS: {results['nss_note']}")

    _section("7. LLM_SSL_VERIFY env")
    raw = LLM_SSL_VERIFY.strip().lower()
    disabled = raw in ("false", "0", "no", "off")
    results["llm_ssl_verify_disabled"] = disabled
    print(f"  LLM_SSL_VERIFY={raw}  {'DISABLED' if disabled else 'ENABLED'}")
    if not disabled:
        try:
            ctx = _ctx_cafile()
            results["ssl_context_works"] = True
            print(f"  SSLContext(cafile) works, verify_mode={ctx.verify_mode}")
        except Exception as e:
            results["ssl_context_works"] = False
            results["ssl_context_error"] = str(e)
            print(f"  SSLContext(cafile) FAIL: {e}")

    _section("8. AI Diagnosis Summary")
    summary = {}
    for k in [
        "target_url", "host", "llm_ssl_verify", "llm_ssl_verify_disabled",
        "ca_bundle_certs", "ca_hash_links", "custom_certs_count",
        "openssl_default", "openssl_cafile", "openssl_capath", "openssl_chain",
        "httpx_true", "httpx_cafile", "httpx_capath", "httpx_cafile_capath", "httpx_false",
        "requests_true", "requests_cafile", "requests_capath", "requests_false",
        "nss_available", "nss_certs_count", "ssl_context_works",
    ]:
        if k in results:
            summary[k] = results[k]
    summary["ai_diagnosis"] = _diagnose(summary)
    print(json.dumps(summary, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()