#!/usr/bin/env python3
"""FinDataHub 数据库状态检测"""
import sys
import json
import os

MYSQL_HOST = os.environ.get("MYSQL_HOST", "localhost")
MYSQL_PORT = int(os.environ.get("MYSQL_PORT", "3306"))
MYSQL_USER = os.environ.get("MYSQL_USER", "root")
MYSQL_PASSWORD = os.environ.get("MYSQL_PASSWORD", "")
MYSQL_DATABASE = os.environ.get("MYSQL_DATABASE", "news_flash")
MYSQL_DATABASE_AKSHARE = os.environ.get("MYSQL_DATABASE_AKSHARE", "akshare_data")
MYSQL_DATABASE_TUSHARE = os.environ.get("MYSQL_DATABASE_TUSHARE", "tushare_data")


def connect(host, port, user, password, database=None):
    import pymysql
    return pymysql.connect(
        host=host, port=port, user=user, password=password,
        database=database, charset="utf8mb4",
        connect_timeout=5
    )


def main():
    result = {
        "mysql_reachable": False,
        "news_flash_exists": False,
        "news_flash_tables": [],
        "akshare_data_exists": False,
        "akshare_data_tables": [],
        "tushare_data_exists": False,
        "tushare_data_tables": [],
        "ready": False,
        "error": None
    }

    try:
        conn = connect(MYSQL_HOST, MYSQL_PORT, MYSQL_USER, MYSQL_PASSWORD)
        result["mysql_reachable"] = True

        cursor = conn.cursor()

        cursor.execute("SHOW DATABASES")
        dbs = [row[0] for row in cursor.fetchall()]

        result["news_flash_exists"] = MYSQL_DATABASE in dbs
        result["akshare_data_exists"] = MYSQL_DATABASE_AKSHARE in dbs
        result["tushare_data_exists"] = MYSQL_DATABASE_TUSHARE in dbs

        if MYSQL_DATABASE in dbs:
            conn.select_db(MYSQL_DATABASE)
            cursor.execute("SHOW TABLES")
            result["news_flash_tables"] = [row[0] for row in cursor.fetchall()]

        if MYSQL_DATABASE_AKSHARE in dbs:
            conn.select_db(MYSQL_DATABASE_AKSHARE)
            cursor.execute("SHOW TABLES")
            result["akshare_data_tables"] = [row[0] for row in cursor.fetchall()]

        if MYSQL_DATABASE_TUSHARE in dbs:
            conn.select_db(MYSQL_DATABASE_TUSHARE)
            cursor.execute("SHOW TABLES")
            result["tushare_data_tables"] = [row[0] for row in cursor.fetchall()]

        conn.close()

        required_nf_tables = ["news_sources", "news_flash", "push_log", "fetch_log"]
        has_nf_tables = all(t in result["news_flash_tables"] for t in required_nf_tables)

        result["ready"] = result["news_flash_exists"] and has_nf_tables

    except Exception as e:
        result["error"] = str(e)

    print(json.dumps(result, ensure_ascii=False))


if __name__ == "__main__":
    main()