Medical GraphRAG Assistant

"""
Health check functions for AWS GPU NIM RAG system validation.

Provides structured health checks for all system components:
- GPU availability and utilization
- Docker GPU runtime
- IRIS database connectivity and schema
- NIM LLM service health and inference

All functions return structured results with pass/fail status
and diagnostic messages for troubleshooting.
"""

import subprocess
import json
import urllib.request
import urllib.error
from typing import Dict, List, Any, Optional
from dataclasses import dataclass, asdict


@dataclass
class HealthCheckResult:
    """Structured result from a health check.

    Attributes:
        component: Component name being checked
        status: 'pass' or 'fail'
        message: Human-readable status message
        details: Optional additional diagnostic information
    """
    component: str
    status: str  # 'pass' or 'fail'
    message: str
    details: Optional[Dict[str, Any]] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary representation."""
        return asdict(self)


def gpu_check() -> HealthCheckResult:
    """
    Check GPU availability and basic functionality.

    Verifies:
    - nvidia-smi command is available
    - GPU device is detected
    - Driver and CUDA versions

    Returns:
        HealthCheckResult with GPU information or error details
    """
    try:
        # Check if nvidia-smi exists
        result = subprocess.run(
            ['nvidia-smi', '--query-gpu=name,driver_version,memory.total', '--format=csv,noheader'],
            capture_output=True,
            text=True,
            timeout=10
        )

        if result.returncode != 0:
            return HealthCheckResult(
                component="GPU",
                status="fail",
                message="GPU not accessible",
                details={"error": result.stderr}
            )

        # Parse GPU information
        gpu_info = result.stdout.strip().split(', ')

        if len(gpu_info) < 3:
            return HealthCheckResult(
                component="GPU",
                status="fail",
                message="Could not parse GPU information",
                details={"output": result.stdout}
            )

        gpu_name = gpu_info[0]
        driver_version = gpu_info[1]
        memory_mb = gpu_info[2].split()[0]

        # Get CUDA version
        cuda_result = subprocess.run(
            ['nvidia-smi'],
            capture_output=True,
            text=True,
            timeout=10
        )

        cuda_version = "Unknown"
        for line in cuda_result.stdout.split('\n'):
            if "CUDA Version" in line:
                parts = line.split("CUDA Version:")
                if len(parts) > 1:
                    cuda_version = parts[1].strip().split()[0]
                break

        return HealthCheckResult(
            component="GPU",
            status="pass",
            message=f"GPU detected: {gpu_name}",
            details={
                "gpu_name": gpu_name,
                "driver_version": driver_version,
                "memory_mb": memory_mb,
                "cuda_version": cuda_version
            }
        )

    except FileNotFoundError:
        return HealthCheckResult(
            component="GPU",
            status="fail",
            message="nvidia-smi not found - GPU drivers may not be installed",
            details={"suggestion": "Run: ./scripts/aws/install-gpu-drivers.sh"}
        )
    except subprocess.TimeoutExpired:
        return HealthCheckResult(
            component="GPU",
            status="fail",
            message="nvidia-smi command timed out",
            details={}
        )
    except Exception as e:
        return HealthCheckResult(
            component="GPU",
            status="fail",
            message=f"Unexpected error checking GPU: {str(e)}",
            details={"error_type": type(e).__name__}
        )


def gpu_utilization_check() -> HealthCheckResult:
    """
    Check GPU utilization and memory usage.

    Provides real-time GPU metrics useful for monitoring workloads.

    Returns:
        HealthCheckResult with utilization metrics
    """
    try:
        result = subprocess.run(
            ['nvidia-smi', '--query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu',
             '--format=csv,noheader,nounits'],
            capture_output=True,
            text=True,
            timeout=10
        )

        if result.returncode != 0:
            return HealthCheckResult(
                component="GPU Utilization",
                status="fail",
                message="Could not query GPU utilization",
                details={"error": result.stderr}
            )

        # Parse utilization info
        util_info = result.stdout.strip().split(', ')

        if len(util_info) < 4:
            return HealthCheckResult(
                component="GPU Utilization",
                status="fail",
                message="Could not parse GPU utilization",
                details={"output": result.stdout}
            )

        gpu_util_pct = float(util_info[0])
        memory_used_mb = float(util_info[1])
        memory_total_mb = float(util_info[2])
        temperature_c = float(util_info[3])

        memory_util_pct = (memory_used_mb / memory_total_mb) * 100 if memory_total_mb > 0 else 0

        return HealthCheckResult(
            component="GPU Utilization",
            status="pass",
            message=f"GPU utilization: {gpu_util_pct}%",
            details={
                "gpu_utilization_pct": gpu_util_pct,
                "memory_used_mb": memory_used_mb,
                "memory_total_mb": memory_total_mb,
                "memory_utilization_pct": round(memory_util_pct, 1),
                "temperature_c": temperature_c
            }
        )

    except Exception as e:
        return HealthCheckResult(
            component="GPU Utilization",
            status="fail",
            message=f"Error checking GPU utilization: {str(e)}",
            details={"error_type": type(e).__name__}
        )


def docker_gpu_check() -> HealthCheckResult:
    """
    Check Docker GPU runtime configuration.

    Verifies:
    - Docker is installed
    - Docker can access GPU via --gpus flag
    - GPU is accessible inside containers

    Returns:
        HealthCheckResult with Docker GPU status
    """
    try:
        # Test GPU access in container
        result = subprocess.run(
            ['docker', 'run', '--rm', '--gpus', 'all',
             'nvidia/cuda:12.2.0-base-ubuntu22.04', 'nvidia-smi'],
            capture_output=True,
            text=True,
            timeout=60
        )

        if result.returncode != 0:
            return HealthCheckResult(
                component="Docker GPU Runtime",
                status="fail",
                message="Docker cannot access GPU",
                details={
                    "error": result.stderr,
                    "suggestion": "Run: ./scripts/aws/setup-docker-gpu.sh"
                }
            )

        return HealthCheckResult(
            component="Docker GPU Runtime",
            status="pass",
            message="Docker can access GPU",
            details={}
        )

    except FileNotFoundError:
        return HealthCheckResult(
            component="Docker GPU Runtime",
            status="fail",
            message="Docker not found",
            details={"suggestion": "Install Docker first"}
        )
    except subprocess.TimeoutExpired:
        return HealthCheckResult(
            component="Docker GPU Runtime",
            status="fail",
            message="Docker GPU test timed out",
            details={}
        )
    except Exception as e:
        return HealthCheckResult(
            component="Docker GPU Runtime",
            status="fail",
            message=f"Error checking Docker GPU: {str(e)}",
            details={"error_type": type(e).__name__}
        )


def iris_connection_check(host: str = "localhost", port: int = 1972,
                          namespace: str = "DEMO", username: str = "_SYSTEM",
                          password: str = "SYS") -> HealthCheckResult:
    """
    Check IRIS database connectivity.

    Args:
        host: IRIS host address
        port: IRIS SuperServer port
        namespace: IRIS namespace
        username: Database username
        password: Database password

    Returns:
        HealthCheckResult with connection status
    """
    try:
        import iris

        # Attempt connection
        conn = iris.connect(host, port, namespace, username, password)
        cursor = conn.cursor()

        # Test with simple query
        cursor.execute("SELECT 1")
        result = cursor.fetchone()

        conn.close()

        if result and result[0] == 1:
            return HealthCheckResult(
                component="IRIS Connection",
                status="pass",
                message=f"Connected to IRIS at {host}:{port}/{namespace}",
                details={
                    "host": host,
                    "port": port,
                    "namespace": namespace
                }
            )
        else:
            return HealthCheckResult(
                component="IRIS Connection",
                status="fail",
                message="Query test failed",
                details={"result": str(result)}
            )

    except ImportError:
        return HealthCheckResult(
            component="IRIS Connection",
            status="fail",
            message="iris Python module not installed",
            details={"suggestion": "pip install intersystems-iris"}
        )
    except Exception as e:
        return HealthCheckResult(
            component="IRIS Connection",
            status="fail",
            message=f"Connection failed: {str(e)}",
            details={
                "error_type": type(e).__name__,
                "host": host,
                "port": port,
                "suggestion": "Check IRIS container is running: docker ps | grep iris"
            }
        )


def iris_tables_check(host: str = "localhost", port: int = 1972,
                      namespace: str = "DEMO", username: str = "_SYSTEM",
                      password: str = "SYS") -> HealthCheckResult:
    """
    Check IRIS vector table schema.

    Verifies:
    - ClinicalNoteVectors table exists
    - MedicalImageVectors table exists

    Args:
        host: IRIS host address
        port: IRIS SuperServer port
        namespace: IRIS namespace
        username: Database username
        password: Database password

    Returns:
        HealthCheckResult with table existence status
    """
    try:
        import iris

        conn = iris.connect(host, port, namespace, username, password)
        cursor = conn.cursor()

        # Check ClinicalNoteVectors table
        cursor.execute("""
            SELECT COUNT(*)
            FROM INFORMATION_SCHEMA.TABLES
            WHERE TABLE_SCHEMA=? AND TABLE_NAME='ClinicalNoteVectors'
        """, (namespace,))
        clinical_exists = cursor.fetchone()[0] > 0

        # Check MedicalImageVectors table
        cursor.execute("""
            SELECT COUNT(*)
            FROM INFORMATION_SCHEMA.TABLES
            WHERE TABLE_SCHEMA=? AND TABLE_NAME='MedicalImageVectors'
        """, (namespace,))
        image_exists = cursor.fetchone()[0] > 0

        conn.close()

        tables_found = []
        tables_missing = []

        if clinical_exists:
            tables_found.append("ClinicalNoteVectors")
        else:
            tables_missing.append("ClinicalNoteVectors")

        if image_exists:
            tables_found.append("MedicalImageVectors")
        else:
            tables_missing.append("MedicalImageVectors")

        if len(tables_found) == 2:
            return HealthCheckResult(
                component="IRIS Tables",
                status="pass",
                message="All vector tables exist",
                details={
                    "tables_found": tables_found,
                    "namespace": namespace
                }
            )
        elif len(tables_found) > 0:
            return HealthCheckResult(
                component="IRIS Tables",
                status="fail",
                message="Some vector tables missing",
                details={
                    "tables_found": tables_found,
                    "tables_missing": tables_missing,
                    "suggestion": "Run schema creation scripts"
                }
            )
        else:
            return HealthCheckResult(
                component="IRIS Tables",
                status="fail",
                message="No vector tables found",
                details={
                    "tables_missing": tables_missing,
                    "suggestion": "Run: python src/setup/create_text_vector_table.py"
                }
            )

    except ImportError:
        return HealthCheckResult(
            component="IRIS Tables",
            status="fail",
            message="iris Python module not installed",
            details={"suggestion": "pip install intersystems-iris"}
        )
    except Exception as e:
        return HealthCheckResult(
            component="IRIS Tables",
            status="fail",
            message=f"Table check failed: {str(e)}",
            details={"error_type": type(e).__name__}
        )


def nim_llm_health_check(host: str = "localhost", port: int = 8001) -> HealthCheckResult:
    """
    Check NIM LLM service health endpoint.

    Args:
        host: NIM service host
        port: NIM service port

    Returns:
        HealthCheckResult with service health status
    """
    try:
        url = f"http://{host}:{port}/health"

        req = urllib.request.Request(url)

        with urllib.request.urlopen(req, timeout=10) as response:
            if response.status == 200:
                return HealthCheckResult(
                    component="NIM LLM Health",
                    status="pass",
                    message=f"NIM LLM service healthy at {host}:{port}",
                    details={"endpoint": url}
                )
            else:
                return HealthCheckResult(
                    component="NIM LLM Health",
                    status="fail",
                    message=f"Health endpoint returned status {response.status}",
                    details={"status_code": response.status}
                )

    except urllib.error.URLError as e:
        return HealthCheckResult(
            component="NIM LLM Health",
            status="fail",
            message="Health endpoint not accessible",
            details={
                "error": str(e),
                "endpoint": f"http://{host}:{port}/health",
                "suggestion": "NIM may still be initializing - check: docker logs nim-llm"
            }
        )
    except Exception as e:
        return HealthCheckResult(
            component="NIM LLM Health",
            status="fail",
            message=f"Error checking health: {str(e)}",
            details={"error_type": type(e).__name__}
        )


def nim_llm_inference_test(host: str = "localhost", port: int = 8001) -> HealthCheckResult:
    """
    Test NIM LLM inference with a simple query.

    Args:
        host: NIM service host
        port: NIM service port

    Returns:
        HealthCheckResult with inference test status and response
    """
    try:
        url = f"http://{host}:{port}/v1/chat/completions"

        payload = {
            "model": "meta/llama-3.1-8b-instruct",
            "messages": [
                {"role": "user", "content": "What is 2+2? Answer with just the number."}
            ],
            "max_tokens": 10
        }

        data = json.dumps(payload).encode('utf-8')
        req = urllib.request.Request(
            url,
            data=data,
            headers={'Content-Type': 'application/json'}
        )

        with urllib.request.urlopen(req, timeout=30) as response:
            if response.status == 200:
                response_data = json.loads(response.read().decode('utf-8'))

                if 'choices' in response_data and len(response_data['choices']) > 0:
                    answer = response_data['choices'][0]['message']['content']

                    return HealthCheckResult(
                        component="NIM LLM Inference",
                        status="pass",
                        message="Inference test successful",
                        details={
                            "test_query": "What is 2+2?",
                            "response": answer.strip(),
                            "endpoint": url
                        }
                    )
                else:
                    return HealthCheckResult(
                        component="NIM LLM Inference",
                        status="fail",
                        message="Response missing expected fields",
                        details={"response": response_data}
                    )
            else:
                return HealthCheckResult(
                    component="NIM LLM Inference",
                    status="fail",
                    message=f"Inference endpoint returned status {response.status}",
                    details={"status_code": response.status}
                )

    except urllib.error.URLError as e:
        return HealthCheckResult(
            component="NIM LLM Inference",
            status="fail",
            message="Inference endpoint not accessible",
            details={
                "error": str(e),
                "endpoint": url,
                "suggestion": "Model may still be loading - check: docker logs nim-llm"
            }
        )
    except json.JSONDecodeError as e:
        return HealthCheckResult(
            component="NIM LLM Inference",
            status="fail",
            message="Could not parse response JSON",
            details={"error": str(e)}
        )
    except Exception as e:
        return HealthCheckResult(
            component="NIM LLM Inference",
            status="fail",
            message=f"Inference test failed: {str(e)}",
            details={"error_type": type(e).__name__}
        )


def run_all_checks(iris_host: str = "localhost", iris_port: int = 1972,
                  nim_host: str = "localhost", nim_port: int = 8001,
                  skip_gpu: bool = False, skip_docker: bool = False,
                  skip_iris: bool = False, skip_nim: bool = False) -> List[HealthCheckResult]:
    """
    Run all health checks and return results.

    Args:
        iris_host: IRIS database host
        iris_port: IRIS database port
        nim_host: NIM service host
        nim_port: NIM service port
        skip_gpu: Skip GPU checks
        skip_docker: Skip Docker checks
        skip_iris: Skip IRIS checks
        skip_nim: Skip NIM checks

    Returns:
        List of HealthCheckResult objects
    """
    results = []

    if not skip_gpu:
        results.append(gpu_check())
        results.append(gpu_utilization_check())

    if not skip_docker:
        results.append(docker_gpu_check())

    if not skip_iris:
        results.append(iris_connection_check(iris_host, iris_port))
        results.append(iris_tables_check(iris_host, iris_port))

    if not skip_nim:
        results.append(nim_llm_health_check(nim_host, nim_port))
        results.append(nim_llm_inference_test(nim_host, nim_port))

    return results


if __name__ == "__main__":
    """Run all health checks when executed as script."""
    print("Running health checks...")
    print()

    results = run_all_checks()

    passed = sum(1 for r in results if r.status == "pass")
    failed = sum(1 for r in results if r.status == "fail")

    for result in results:
        status_symbol = "✓" if result.status == "pass" else "✗"
        print(f"{status_symbol} {result.component}: {result.message}")

        if result.details and result.status == "fail":
            if "suggestion" in result.details:
                print(f"   Suggestion: {result.details['suggestion']}")
            if "error" in result.details:
                print(f"   Error: {result.details['error']}")
        print()

    print(f"Results: {passed} passed, {failed} failed")

    # Exit with error code if any check failed
    exit(0 if failed == 0 else 1)