import time
from time import sleep
from prometheus_client import start_http_server, Gauge
import argparse
import pymysql
import requests


def service_status():
    # 数据库连接配置
    config = {
        'host': args.db_ip,
        'port': int(args.db_port),
        'database': 'monitor',  # 替换为你的数据库名称
        'user': args.db_user,  # 替换为你的数据库用户名
        'password': args.db_password  # 替换为你的数据库密码
    }

    # 连接到数据库
    try:
        connection = pymysql.connect(**config)
        if connection.open:
            # 创建一个游标对象
            cursor = connection.cursor()

            # 执行一个查询来获取项目和服务的相关信息
            query = """
                    SELECT a.project,a.service_name,a.url,b.owner_name,b.owner_email 
                    from monitor.service a JOIN monitor.project b ON a.project=b.project 
                    order by a.project ,a.service_name; 
            """

            cursor.execute(query)

            # 获取所有记录
            records = cursor.fetchall()

            # 将记录转换为字典列表
            result = []
            for record in records:
                try:
                    status = requests.get(record[2]).status_code
                except Exception:
                    # 当网页无法访问时，将状态码设置为777
                    status = 777
                result.append({
                    'project': record[0],
                    'service_name': record[1],
                    'url': record[2],
                    'status': status,
                    'owner_name': record[3],
                    'owner_email': record[4],
                    'chk_time': time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
                })

            return result

    except Exception as e:
        print(f"数据库连接错误: {e}")

    finally:
        # 关闭游标和连接
        if connection.open:
            cursor.close()
            connection.close()


if __name__ == '__main__':

    # 创建一个Gauge指标
    metric = Gauge('zzsn_service_status', 'service running status code',
                   ['project', 'service_name', 'url', 'owner_name', 'owner_email', 'chk_time'])

    # 解析命令行参数
    parser = argparse.ArgumentParser(description='这是一个服务状态检查程序，说明如下：')
    parser.add_argument('--port', '-p', default='8761', help='监听端口，8761')
    parser.add_argument('--db_ip', default='114.116.11.225', help='指定mysql的ip，默认为114.116.11.225')
    parser.add_argument('--db_port', default='3306', help='指定mysql端口，默认为3306')
    parser.add_argument('--db_user', default='monitor', help='指定mysql的用户名，monitor')
    parser.add_argument('--db_password', default='@ZZsn9988', help='指定mysql的密码,默认值你猜')
    args = parser.parse_args()
    # 启动HTTP服务器，监听指定的端口，并添加基本身份验证
    start_http_server(int(args.port))
    print("zzsn Service Exporter for Prometheus is running...")
    print(f"System args is \n"
          f"Listen port:{args.port}\n"
          f"Mysql IP:{args.db_ip}\n"
          f"Mysql port:{args.db_port}\n"
          f"Mysql username:{args.db_user}\n"
          f"Mysql Database:monitor")
    while True:
        # 更新指标的值
        res = service_status()
        metric.clear()
        for info in res:
            metric.labels(project=info['project'], service_name=info['service_name'],
                          url=info['url'], owner_name=info['owner_name'],
                          owner_email=info['owner_email'], chk_time=info['chk_time']).set(info['status'])

        # 每隔30秒更新一次指标的值
        sleep(30)
