from flask import Flask
from prometheus_client import Gauge, generate_latest, CONTENT_TYPE_LATEST
import subprocess
import time
import os
import re
import threading

app = Flask(__name__)

pv_usage_bytes = Gauge('nfs_pv_usage_bytes', 'PV usage in bytes', ['pv_name', 'request_name'])
pv_last_check = Gauge('nfs_pv_last_check_timestamp', 'Last check timestamp', ['pv_name', 'request_name'])

PVC_UUID_RE = re.compile(r'(pvc-[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})$')

def extract_pv_name(dir_name):
    m = PVC_UUID_RE.search(dir_name)
    return m.group(1) if m else dir_name

cache = {}
CACHE_TTL = 300
du_lock = threading.Lock()  # 동시 du 방지

def get_pv_usage(pv_name):
    now = time.time()
    
    # 캐시 확인 (락 없이)
    if pv_name in cache and now - cache[pv_name]['time'] < CACHE_TTL:
        return cache[pv_name]['bytes']
    
    # du는 한 번에 하나만
    with du_lock:
        # 락 획득 후 다시 캐시 확인
        if pv_name in cache and time.time() - cache[pv_name]['time'] < CACHE_TTL:
            return cache[pv_name]['bytes']
        
        path = f"/mnt/pvs/{pv_name}"
        if not os.path.exists(path):
            return None
        
        result = subprocess.run(
            ['ionice', '-c', '3', 'nice', '-n', '19', 'du', '-sb', path],
            capture_output=True, text=True, timeout=300
        )
        
        if result.returncode != 0:
            return None
        
        usage = int(result.stdout.split()[0])
        real_pv = extract_pv_name(pv_name)
        cache[pv_name] = {'bytes': usage, 'time': time.time()}
        pv_usage_bytes.labels(pv_name=real_pv, request_name=pv_name).set(usage)
        pv_last_check.labels(pv_name=real_pv, request_name=pv_name).set(time.time())
        
        return usage

@app.route('/usage/<pv_name>')
def usage(pv_name):
    result = get_pv_usage(pv_name)
    if result is None:
        return {'error': 'PV not found or timeout'}, 404
    return {'pv_name': extract_pv_name(pv_name), 'request_name': pv_name, 'bytes': result}

@app.route('/metrics')
def metrics():
    return generate_latest(), 200, {'Content-Type': CONTENT_TYPE_LATEST}

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8080, threaded=True)