import sqlite3
from flask import Flask, Response

app = Flask(__name__)
DB_FILE = "vms.db"

def get_or_assign_hostname(mac):
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    
    # 1. Create table if it doesn't exist
    cursor.execute('''CREATE TABLE IF NOT EXISTS vms 
                      (id INTEGER PRIMARY KEY AUTOINCREMENT, mac TEXT UNIQUE, hostname TEXT, used INTEGER)''')
    
    # 2. Check if this MAC address already has a name assigned
    cursor.execute("SELECT hostname FROM vms WHERE mac = ?", (mac,))
    row = cursor.fetchone()
    if row:
        conn.close()
        return row[0]
    
    # 3. Find the lowest unused VM number
    # First, let's seed the database with names if it's completely empty
    cursor.execute("SELECT COUNT(*) FROM vms")
    if cursor.fetchone()[0] == 0:
        for i in range(10000):
            cursor.execute("INSERT INTO vms (hostname, used) VALUES (?, 0)", (f"vm{i}",))
        conn.commit()
    
    # Find the first available name and lock it to this MAC
    cursor.execute("SELECT id, hostname FROM vms WHERE used = 0 ORDER BY id LIMIT 1")
    avail = cursor.fetchone()
    
    if avail:
        vm_id, hostname = avail
        cursor.execute("UPDATE vms SET mac = ?, used = 1 WHERE id = ?", (mac, vm_id))
        conn.commit()
        conn.close()
        return hostname
    
    conn.close()
    return "vm-overflow"

@app.route('/boot/<mac>')
def generate_ipxe(mac):
    hostname = get_or_assign_hostname(mac)
    
    # Dynamically generate the iPXE script with the unique hostname injected!
    ipxe_script = f"""#!ipxe
set server_url http://192.168.88.16
initrd ${{server_url}}/initrd.gz
initrd ${{server_url}}/preseed.cfg preseed.cfg
kernel ${{server_url}}/vmlinuz initrd=initrd.gz auto=true priority=critical DEBCONF_DEBUG=5 netcfg/get_hostname={hostname} netcfg/get_domain=lab.baldwinit.org
boot
"""
    return Response(ipxe_script, mimetype='text/plain')

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