from flask import Flask, jsonify
import subprocess
import os
import socket
import psutil


hostname = socket.gethostname()

app = Flask(__name__)

# refresh the text file
def refresh_crystal_info():
    try:
        # Run CrystalDiskInfo command to generate DiskInfo.txt
        return subprocess.run(["DiskInfo64.exe", "/CopyExit"])
    except Exception as e:
        return {"error": str(e)}

def get_crystal_disk_info():
    try:
        # Read the generated DiskInfo.txt file
        with open("DiskInfo.txt", "r") as file:
            output = file.read()

        # Initialize a list to hold each drive's data
        drives = []

        # Split the file content into sections for each drive
        drive_sections = output.split('----------------------------------------------------------------------------')
        for section in drive_sections:
            lines = section.strip().splitlines()
            data = {
                "Hostname": None,
                "Disk Size": None,
                "Model": None,
                "Serial Number": None,
                "Firmware": None,
                "Temperature": None,
                "Health Status": None,
                "Power On Hours": None,
                "Power On Count": None,
                "Host Writes": None,
                "Wear Level Count": None,
                "Drive Letter": None
            }
            for line in lines: 
                if "Model" in line:
                    if ":" in line:
                        data["Model"] = line.split(":", 1)[1].strip()
                elif "Serial Number" in line:
                    if ":" in line:
                        data["Serial Number"] = line.split(":", 1)[1].strip()
                elif "Firmware" in line:
                    if ":" in line:
                        data["Firmware"] = line.split(":", 1)[1].strip()
                elif "Temperature" in line:
                    if ":" in line:
                        data["Temperature"] = line.split(":", 1)[1].strip()
                elif "Health Status" in line:
                    if ":" in line:
                        data["Health Status"] = line.split(":", 1)[1].strip()
                elif "Power On Hours" in line:
                    if ":" in line:
                        data["Power On Hours"] = line.split(":", 1)[1].strip()
                elif "Power On Count" in line:
                    if ":" in line:
                        data["Power On Count"] = line.split(":", 1)[1].strip()
                elif "Host Writes" in line:
                    if ":" in line:
                        data["Host Writes"] = line.split(":", 1)[1].strip()
                elif "Wear Level Count" in line:
                    if ":" in line:
                        data["Wear Level Count"] = line.split(":", 1)[1].strip()
                elif "Drive Letter" in line:
                    if ":" in line:
                        data["Drive Letter"] = line.split(":", 1)[1].strip()
                elif "Disk Size" in line:
                    if ":" in line:
                        data["Disk Size"] = line.split(":", 1)[1].strip()
            
            if any(value is not None for value in data.values()):
                drives.append(data)
                data["Hostname"] =  hostname

        if not drives:
            raise ValueError("No drive data found")

        return {"drives": drives}
    except Exception as e:
        refresh_crystal_info()
        return {"error": str(e)}

def bytes_to_human_readable(bytes):
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if bytes < 1024.0:
            return f"{bytes:.2f} {unit}"
        bytes /= 1024.0

@app.route('/udpate_cache', methods=['GET'])
def udpate_cache():
    return jsonify(refresh_crystal_info())

@app.route('/drive_health', methods=['GET'])
def drive_health():
    return jsonify(get_crystal_disk_info())

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)