import paramiko
import sys
import time

sys.stdout.reconfigure(encoding='utf-8')
key_path = r'C:\Users\Snarf\.ssh\id_ed25519_snarfnet'

def main():
    print("=== Connecting to Shipyard Swarm Manager (140.44.4.71) for Heimdall Remediation ===")
    c = paramiko.SSHClient()
    c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    c.connect("140.44.4.71", username="root", key_filename=key_path, timeout=10)

    def run(cmd):
        print(f"\n>>> {cmd}")
        _, o, e = c.exec_command(cmd)
        res = o.read().decode().strip()
        err = e.read().decode().strip()
        if res:
            print(res)
        if err:
            print("[stderr]", err)
        return res

    # 1. Ensure backup directory exists & take safety backup
    run("mkdir -p /mnt/nfs_shares/backups")
    run("cp -v /mnt/nfs_shares/heimdall/www/app.sqlite /mnt/nfs_shares/backups/heimdall_pre_remediation.sqlite")

    # 2. Database repair script executed locally on node
    db_clean_script = '''
import sqlite3
import shutil

db_path = "/mnt/nfs_shares/heimdall/www/app.sqlite"
con = sqlite3.connect(db_path)
cur = con.cursor()

print("Initial total items:", cur.execute("SELECT count(*) FROM items").fetchone()[0])

# Delete all soft-deleted items (deleted_at IS NOT NULL)
cur.execute("DELETE FROM items WHERE deleted_at IS NOT NULL")
print("Deleted soft-deleted items. Remaining:", cur.execute("SELECT count(*) FROM items").fetchone()[0])

# Remove legacy Pi-hole items (Pi-hole 1 and Pi-hole 2) since Pi-hole stack is decommissioned
cur.execute("DELETE FROM items WHERE title LIKE 'Pi-hole%'")
print("Purged legacy Pi-hole items. Remaining:", cur.execute("SELECT count(*) FROM items").fetchone()[0])

# Deduplicate items by title just in case, keeping the highest ID
cur.execute("""
DELETE FROM items 
WHERE id NOT IN (
    SELECT MAX(id) 
    FROM items 
    GROUP BY title
) AND id != 0
""")

# Re-sequence the 'order' field sequentially
rows = cur.execute("SELECT id FROM items WHERE id != 0 ORDER BY [order] ASC, title ASC").fetchall()
for idx, (item_id,) in enumerate(rows):
    cur.execute("UPDATE items SET [order] = ? WHERE id = ?", (idx, item_id))

con.commit()
cur.execute("REINDEX")
cur.execute("VACUUM")
con.close()

print("Database sanitized, re-indexed, and vacuumed successfully!")
'''
    sftp = c.open_sftp()
    with sftp.file('/tmp/clean_heimdall_db.py', 'w') as f:
        f.write(db_clean_script)
    sftp.close()

    run("python3 /tmp/clean_heimdall_db.py")

    # 3. Create local high-performance storage directory on snshpd4lnxswm01
    run("mkdir -p /var/lib/docker-data/heimdall/config")
    run("rsync -av --progress /mnt/nfs_shares/heimdall/ /var/lib/docker-data/heimdall/config/")
    run("chown -R 1000:1000 /var/lib/docker-data/heimdall/config")
    run("ls -la /var/lib/docker-data/heimdall/config/www/")

    # 4. Generate updated Docker Compose file with node pinning & local volume
    new_compose = '''version: "3.8"

services:
  heimdall:
    image: lscr.io/linuxserver/heimdall:latest
    ports:
      - target: 80
        published: 85
        mode: ingress
      - target: 443
        published: 403
        mode: ingress
    environment:
      - PUID=1000
      - PGID=1000
      - TZ=America/New_York
    volumes:
      - type: bind
        source: /var/lib/docker-data/heimdall/config
        target: /config
    deploy:
      mode: replicated
      replicas: 1
      placement:
        constraints:
          - node.hostname == snshpd4lnxswm01
      restart_policy:
        condition: on-failure

networks:
  default:
    name: homelab_public
    external: true
'''
    sftp = c.open_sftp()
    with sftp.file('/tmp/heimdall_compose.yml', 'w') as f:
        f.write(new_compose)
    sftp.close()

    # 5. Redeploy Heimdall stack
    run("docker stack deploy -c /tmp/heimdall_compose.yml heimdall")
    time.sleep(5)
    run("docker stack ps heimdall")

    # 6. Setup daily atomic SQLite backup cron job
    cron_cmd = '(crontab -l 2>/dev/null | grep -v "heimdall_backup"; echo "30 3 * * * python3 -c \\"import sqlite3; con=sqlite3.connect(\'/var/lib/docker-data/heimdall/config/www/app.sqlite\'); b=sqlite3.connect(\'/mnt/nfs_shares/backups/heimdall_daily.sqlite\'); con.backup(b); b.close(); con.close()\\" >/dev/null 2>&1") | crontab -'
    run(cron_cmd)
    run("crontab -l | grep heimdall")

    c.close()
    print("\n=== Heimdall Storage & DB Remediation Complete ===")

if __name__ == '__main__':
    main()