#!/usr/bin/env python3

import requests
import time
import os
import json
from datetime import datetime

# --- Configuration for AIVA's Elestio Connection ---
# IMPORTANT: For production, ensure ELESTIO_API_URL and ELESTIO_API_KEY
# are set as environment variables. Example:
# export ELESTIO_API_URL="https://<your-elestio-domain>/v1/chat/completions"
# export ELESTIO_API_KEY="sk-elestio-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"

REPORT_FILE_PATH = "/mnt/e/genesis-system/AIVA/status/elestio_health_report.md"

def test_elestio_connection():
    """
    Tests the connection health to the Elestio/Qwen endpoint.
    Reports latency, availability, and model status.
    """
    api_url = os.getenv("ELESTIO_API_URL")
    api_key = os.getenv("ELESTIO_API_KEY")

    report_data = {
        "timestamp": datetime.now().isoformat(),
        "endpoint": api_url if api_url else "NOT_CONFIGURED",
        "latency_ms": "N/A",
        "availability": "UNKNOWN",
        "http_status_code": "N/A",
        "model_status": "UNKNOWN",
        "error_message": None
    }

    if not api_url or not api_key:
        report_data["availability"] = "FAILED - Configuration Missing"
        report_data["model_status"] = "FAILED - API URL or Key not set in environment variables."
        report_data["error_message"] = "Please set ELESTIO_API_URL and ELESTIO_API_KEY environment variables."
        print("ERROR: Elestio API URL or Key not found in environment variables.")
        return report_data

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    # A simple chat completion request to test the Qwen model
    payload = {
        "model": "qwen-turbo",  # Or specific Qwen model hosted by Elestio, e.g., "qwen-7b-chat"
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "Ping"}
        ],
        "temperature": 0.0, # Keep it deterministic for health checks
        "max_tokens": 5     # Small response to minimize latency impact
    }

    try:
        start_time = time.perf_counter()
        response = requests.post(api_url, headers=headers, json=payload, timeout=10)
        end_time = time.perf_counter()

        latency_ms = (end_time - start_time) * 1000
        report_data["latency_ms"] = f"{latency_ms:.2f}"
        report_data["http_status_code"] = response.status_code

        if response.status_code == 200:
            report_data["availability"] = "Online"
            try:
                response_json = response.json()
                if "choices" in response_json and len(response_json["choices"]) > 0:
                    report_data["model_status"] = "Responsive"
                    # Optionally, check for specific content in the response
                    # e.g., if "Ping" response is expected
                    # first_choice_content = response_json["choices"][0]["message"]["content"]
                    # if "pong" in first_choice_content.lower():
                    #     report_data["model_status"] = "Responsive (Pong received)"
                else:
                    report_data["model_status"] = "Model responded, but invalid format or no choices."
                    report_data["error_message"] = f"API response: {response_json}"
            except json.JSONDecodeError:
                report_data["model_status"] = "Model responded, but invalid JSON."
                report_data["error_message"] = f"Raw response: {response.text}"
        elif response.status_code == 401:
            report_data["availability"] = "Offline - Unauthorized"
            report_data["model_status"] = "Failed - Invalid API Key."
            report_data["error_message"] = response.text
        elif response.status_code == 404:
            report_data["availability"] = "Offline - Endpoint Not Found"
            report_data["model_status"] = "Failed - Check API URL."
            report_data["error_message"] = response.text
        else:
            report_data["availability"] = f"Offline - HTTP {response.status_code}"
            report_data["model_status"] = "Failed - Unexpected HTTP status."
            report_data["error_message"] = response.text

    except requests.exceptions.Timeout:
        report_data["availability"] = "Offline - Timeout"
        report_data["model_status"] = "Failed - Request timed out."
        report_data["error_message"] = "The request took too long to respond."
    except requests.exceptions.ConnectionError as e:
        report_data["availability"] = "Offline - Connection Error"
        report_data["model_status"] = "Failed - Could not connect to host."
        report_data["error_message"] = str(e)
    except requests.exceptions.RequestException as e:
        report_data["availability"] = "Offline - Request Error"
        report_data["model_status"] = "Failed - An unknown request error occurred."
        report_data["error_message"] = str(e)
    except Exception as e:
        report_data["availability"] = "Offline - Script Error"
        report_data["model_status"] = "Failed - An unexpected error in the script occurred."
        report_data["error_message"] = str(e)

    return report_data

def generate_report(report_data):
    """
    Generates a Markdown report from the connection test data.
    """
    markdown_content = f"""
# Elestio/Qwen Connection Health Report

**Generated At:** {report_data["timestamp"]}

## Endpoint Information
- **API URL:** `{report_data["endpoint"]}`

## Connection Status
- **Availability:** {report_data["availability"]}
- **HTTP Status Code:** {report_data["http_status_code"]}
- **Latency (ms):** {report_data["latency_ms"]}

## Model Status
- **Qwen Model Responsiveness:** {report_data["model_status"]}

"""

    if report_data["error_message"]:
        markdown_content += f"""
## Details/Error Message
```
{report_data["error_message"]}
```
"""

    try:
        os.makedirs(os.path.dirname(REPORT_FILE_PATH), exist_ok=True)
        with open(REPORT_FILE_PATH, "w", encoding="utf-8") as f:
            f.write(markdown_content)
        print(f"SUCCESS: Elestio health report generated at {REPORT_FILE_PATH}")
    except IOError as e:
        print(f"ERROR: Could not write report file to {REPORT_FILE_PATH}: {e}")


if __name__ == "__main__":
    print("Initiating Elestio connection health check for Queen AIVA...")
    health_status = test_elestio_connection()
    generate_report(health_status)
    print("Elestio connection health check complete.")