from flask import Flask, jsonify
from flask_apscheduler import APScheduler
import psutil
import os
import requests
from requests import RequestException, Response
import json
from subprocess import check_output

app = Flask(__name__)
scheduler = APScheduler()
app.config['JSONIFY_PRETTYPRINT_REGULAR'] = True

# Bits to Bytes etc
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

# Parse cache and return results
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('----------------------------------------------------------------------------')
        disk_id = 0
        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,
                "Interface": None,
                "Transfer Mode": 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:
                        raw = line.split(":", 1)[1].strip()
                        data["Disk Size"] = raw.split('GB')[0].strip() + ' GB'
                elif "Interface" in line:
                    if ":" in line:
                        data["Interface"] = line.split(":", 1)[1].strip()
                elif "Transfer Mode" in line:
                    if ":" in line:
                        data["Transfer Mode"] = line.split(":", 1)[1].strip()
            # This makes sure something was changed, 
            if any(value is not None for value in data.values()):
                data["Disk ID"] = disk_id
                drives.append(data)
                disk_id = disk_id + 1
                #data["Hostname"] =  "{{ hostname_output.stdout_lines[0] }}"

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

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

# Disk Info Function
def get_disk_info():
    disk_info = []
    partitions = psutil.disk_partitions()
    for partition in partitions:
        usage = psutil.disk_usage(partition.mountpoint)
        disk_info.append({
            'device': partition.device.replace('\\\\', '\\').rstrip('\\'),
            #'mountpoint': partition.mountpoint,
            'fstype': partition.fstype,
            'total': bytes_to_human_readable(usage.total),
            'used': bytes_to_human_readable(usage.used),
            'free': bytes_to_human_readable(usage.free),
            'percent': usage.percent
        })
    return disk_info

# Flask endpoints

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

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

def server_reporter():
    #base_url="https://cosmostat.matt-cloud.com"
    base_url="http://10.200.27.20:5001"
    url = f"{base_url}/storage_client_update"
    drives_dict = get_crystal_disk_info()
    data_dict = {
        "hostname": "{{ hostname_output.stdout_lines[0] }}",
        "drives": drives_dict["drives"],
        "API_KEY": "deadbeef",
        "storage_summary": get_disk_info()

    }
    result = []
    try:
        response = requests.post(url, json=data_dict)

        # Raise an exception for non‑2xx status codes
        response.raise_for_status()
        result = response.json()
    except:
        result = {
            "message": "error"
        }

    # Return the JSON payload
    return result

if __name__ == '__main__':

    # disk info Loop Function
    def update_disk_info():
        diskinfo_command = f"{{ storage_api_root }}\\dist\\DiskInfo64.exe /CopyExit"
        result = check_output(diskinfo_command, shell=True)
        print(result)
        server_reporter()
        return result

    # gonna try something wild
    scheduler.add_job(id='update_disk_info',
                    func=update_disk_info,
                    trigger='interval',
                    seconds=10000)

    scheduler.add_job(id='server_reporter',
                    func=server_reporter,
                    trigger='interval',
                    seconds=15)

    scheduler.init_app(app)
    scheduler.start()   
    
    update_disk_info()

    app.run(host='0.0.0.0', port={{ api_service_port }})