#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SQLite 資料庫監控腳本
收集資料庫狀態、大小、查詢效能等指標
"""

import sqlite3
import os
import time
from datetime import datetime
from flask import Flask
from prometheus_client import Gauge, Counter, Histogram, generate_latest, REGISTRY

app = Flask(__name__)

# 定義 Prometheus 指標
db_size_bytes = Gauge('sqlite_database_size_bytes', '資料庫檔案大小（bytes）', ['database'])
db_record_count = Gauge('sqlite_record_count', '資料表記錄總數', ['database', 'table'])
db_query_duration = Histogram('sqlite_query_duration_seconds', '查詢執行時間', ['query_type'])
db_connection_errors = Counter('sqlite_connection_errors_total', '連接錯誤總數')
db_slow_queries = Counter('sqlite_slow_queries_total', '慢查詢總數（>1秒）', ['table'])

# 資料庫路徑
DATABASE_PATH = '/home/ogt/momo_pro_system/data/momo_database.db'

def get_db_size():
    """獲取資料庫檔案大小"""
    try:
        if os.path.exists(DATABASE_PATH):
            size = os.path.getsize(DATABASE_PATH)
            db_size_bytes.labels(database='momo').set(size)
            return size
        return 0
    except Exception as e:
        print(f"Error getting database size: {e}")
        return 0

def get_table_counts():
    """獲取各資料表的記錄數"""
    try:
        conn = sqlite3.connect(DATABASE_PATH, timeout=5)
        cursor = conn.cursor()
        
        # 獲取所有資料表
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'")
        tables = cursor.fetchall()
        
        for (table_name,) in tables:
            try:
                start_time = time.time()
                cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
                count = cursor.fetchone()[0]
                duration = time.time() - start_time
                
                # 記錄指標
                db_record_count.labels(database='momo', table=table_name).set(count)
                db_query_duration.labels(query_type='count').observe(duration)
                
                # 檢測慢查詢
                if duration > 1.0:
                    db_slow_queries.labels(table=table_name).inc()
                    
            except Exception as e:
                print(f"Error counting table {table_name}: {e}")
        
        conn.close()
        
    except Exception as e:
        print(f"Database connection error: {e}")
        db_connection_errors.inc()

def measure_query_performance():
    """測試常見查詢的效能"""
    try:
        conn = sqlite3.connect(DATABASE_PATH, timeout=5)
        cursor = conn.cursor()
        
        # 測試查詢 1：最近銷售數據
        start_time = time.time()
        cursor.execute("SELECT * FROM realtime_sales_monthly ORDER BY 日期 DESC LIMIT 100")
        cursor.fetchall()
        duration = time.time() - start_time
        db_query_duration.labels(query_type='recent_sales').observe(duration)
        
        if duration > 1.0:
            db_slow_queries.labels(table='realtime_sales_monthly').inc()
        
        # 測試查詢 2：商品統計
        start_time = time.time()
        cursor.execute("SELECT 品牌, COUNT(*) FROM realtime_sales_monthly GROUP BY 品牌 LIMIT 50")
        cursor.fetchall()
        duration = time.time() - start_time
        db_query_duration.labels(query_type='brand_stats').observe(duration)
        
        conn.close()
        
    except Exception as e:
        print(f"Query performance test error: {e}")
        db_connection_errors.inc()

@app.route('/metrics')
def metrics():
    """Prometheus metrics endpoint"""
    # 收集最新指標
    get_db_size()
    get_table_counts()
    measure_query_performance()
    
    # 返回 Prometheus 格式的指標
    return generate_latest(REGISTRY)

@app.route('/health')
def health():
    """健康檢查端點"""
    try:
        if os.path.exists(DATABASE_PATH):
            conn = sqlite3.connect(DATABASE_PATH, timeout=2)
            conn.close()
            return 'OK', 200
        return 'Database file not found', 503
    except Exception as e:
        return f'Error: {str(e)}', 503

if __name__ == '__main__':
    print(f"Starting Database Exporter on http://127.0.0.1:9120/metrics")
    app.run(host='127.0.0.1', port=9120)