import os
import sys
import json
import requests
from dotenv import load_dotenv

# Optional: only needed for Gremlin websocket test
try:
    import websocket
    HAS_WEBSOCKET = True
except ImportError:
    HAS_WEBSOCKET = False


def ok(msg):
    print(f"[ OK ] {msg}")


def fail(msg):
    print(f"[FAIL] {msg}")


def load_env():
    load_dotenv()
    ok("Loaded .env file")


def test_http(name, url, path="", method="GET", json_body=None):
    full_url = url.rstrip("/") + path
    try:
        resp = requests.request(
            method,
            full_url,
            json=json_body,
            timeout=5,
        )
        if resp.status_code < 400:
            ok(f"{name} reachable ({resp.status_code}) → {full_url}")
            return True
        else:
            fail(f"{name} error ({resp.status_code}) → {full_url}")
    except Exception as e:
        fail(f"{name} unreachable → {full_url} ({e})")
    return False


def test_gremlin_ws(url):
    if not HAS_WEBSOCKET:
        fail("Gremlin test skipped (websocket-client not installed)")
        return False

    try:
        ws = websocket.create_connection(url, timeout=5)
        ws.close()
        ok(f"Gremlin websocket reachable → {url}")
        return True
    except Exception as e:
        fail(f"Gremlin websocket unreachable → {url} ({e})")
        return False


def main():
    load_env()

    GREMLIN_URL = os.getenv("GREMLIN_URL", "ws://localhost:8182/gremlin")
    ES_URL = os.getenv("ES_URL", "http://localhost:9200")
    ES_INDEX = os.getenv("ES_INDEX", "concepts")
    IPFS_API = os.getenv("IPFS_API", "http://localhost:5001")
    OLLAMA_URL = os.getenv("OLLAMA_URL", "http://localhost:11434")
    OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3.1:8b")
    OLLAMA_EMBED_MODEL = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text")

    print("\n=== Connectivity checks ===\n")

    # Gremlin
    test_gremlin_ws(GREMLIN_URL)

    # Elasticsearch root
    test_http("Elasticsearch", ES_URL)

    # Elasticsearch index existence
    test_http(
        "Elasticsearch index",
        ES_URL,
        path=f"/{ES_INDEX}",
        method="HEAD",
    )

    # IPFS (Kubo)
    test_http(
        "IPFS API",
        IPFS_API,
        path="/api/v0/version",
        method="POST",
    )

    # Ollama base
    test_http(
        "Ollama",
        OLLAMA_URL,
        path="/api/tags",
    )

    # Ollama model availability (best-effort)
    try:
        resp = requests.get(f"{OLLAMA_URL}/api/tags", timeout=5)
        models = [m["name"] for m in resp.json().get("models", [])]

        if OLLAMA_MODEL in models:
            ok(f"Ollama model available → {OLLAMA_MODEL}")
        else:
            fail(f"Ollama model NOT found → {OLLAMA_MODEL}")

        if OLLAMA_EMBED_MODEL in models:
            ok(f"Ollama embed model available → {OLLAMA_EMBED_MODEL}")
        else:
            fail(f"Ollama embed model NOT found → {OLLAMA_EMBED_MODEL}")

    except Exception as e:
        fail(f"Ollama model check failed ({e})")

    print("\n=== Done ===\n")


if __name__ == "__main__":
    main()