#!/usr/bin/env python3
from flask import Flask, render_template, request, jsonify
import yaml
import secrets
import string
import subprocess
from argon2 import PasswordHasher

app = Flask(__name__)

SERVICE_GROUPS = {
    'admins': [
        'traefik.ai-impress.com',
        'pgadmin.ai-impress.com',
        'rabbitmq.ai-impress.com',
        'n8n.ai-impress.com',
        'portainer.ai-impress.com',
        'status.ai-impress.com',
        'supabase.ai-impress.com',
        'wpp.ai-impress.com',
        'social.ai-impress.com',
        'odoo.ai-impress.com',
        'vault.ai-impress.com',
        'chat.ai-impress.com',
        'crm.ai-impress.com',
        'wiki.ai-impress.com',
        'portal.ai-impress.com',
        'llm.ai-impress.com'
    ],
    'developers': [
        'n8n.ai-impress.com',
        'portainer.ai-impress.com',
        'traefik.ai-impress.com',
        'status.ai-impress.com',
        'odoo.ai-impress.com',
        'supabase.ai-impress.com',
        'vault.ai-impress.com',
        'llm.ai-impress.com'
    ],
    'analysts': [
        'pgadmin.ai-impress.com',
        'status.ai-impress.com',
        'odoo.ai-impress.com',
        'crm.ai-impress.com'
    ],
    'users': [
        'wiki.ai-impress.com',
        'status.ai-impress.com',
        'chat.ai-impress.com',
        'crm.ai-impress.com',
        'portal.ai-impress.com'
    ]
}

USERS_FILE = '/config/users_database.yml'

ph = PasswordHasher(
    time_cost=3,
    memory_cost=65536,
    parallelism=4,
    hash_len=32,
    salt_len=16
)

def generate_password(length=16):
    alphabet = string.ascii_letters + string.digits + "!@#$%^&*"
    return ''.join(secrets.choice(alphabet) for _ in range(length))

def hash_password(password):
    return ph.hash(password)

def restart_authelia():
    try:
        subprocess.run(['docker', 'restart', 'authelia'], check=True, capture_output=True)
        return True
    except Exception as e:
        print(f"Failed to restart authelia: {e}")
        return False

def load_users():
    with open(USERS_FILE, 'r') as f:
        return yaml.safe_load(f) or {'users': {}}

def save_users(data):
    with open(USERS_FILE, 'w') as f:
        yaml.dump(data, f, default_flow_style=False, sort_keys=False)

def get_all_services():
    """Extract all unique services from SERVICE_GROUPS"""
    all_services = set()
    for services in SERVICE_GROUPS.values():
        all_services.update(services)
    return sorted(all_services)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/api/groups')
def get_groups():
    return jsonify({
        'groups': list(SERVICE_GROUPS.keys()),
        'services': SERVICE_GROUPS
    })

@app.route('/api/services')
def get_services():
    """Return all available services"""
    return jsonify({
        'services': get_all_services()
    })

@app.route('/api/users')
def get_users():
    data = load_users()
    users = []
    for username, info in data['users'].items():
        # Get services from groups
        group_services = set()
        for group in info.get('groups', []):
            group_services.update(SERVICE_GROUPS.get(group, []))

        # Get individual services
        individual_services = info.get('services', [])

        users.append({
            'username': username,
            'email': info.get('email', ''),
            'displayname': info.get('displayname', ''),
            'groups': info.get('groups', []),
            'group_services': sorted(group_services),
            'individual_services': individual_services,
            'all_services': sorted(set(list(group_services) + individual_services))
        })
    return jsonify(users)

@app.route('/api/users', methods=['POST'])
def add_user():
    try:
        data = request.json
        email = data.get('email')
        displayname = data.get('displayname')
        groups = data.get('groups', [])
        services = data.get('services', [])  # Individual services

        # Use email as username for consistency with Authelia login
        username = email

        users_data = load_users()

        if username in users_data['users']:
            return jsonify({'error': 'User already exists'}), 400

        password = generate_password()
        password_hash = hash_password(password)

        users_data['users'][username] = {
            'displayname': displayname,
            'password': password_hash,
            'email': email,
            'groups': groups,
            'services': services
        }

        save_users(users_data)
        restart_authelia()

        return jsonify({
            'success': True,
            'username': username,
            'password': password
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/users/<username>', methods=['DELETE'])
def delete_user(username):
    try:
        users_data = load_users()

        if username not in users_data['users']:
            return jsonify({'error': 'User not found'}), 404

        del users_data['users'][username]
        save_users(users_data)
        restart_authelia()

        return jsonify({'success': True})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/users/<username>/groups', methods=['PUT'])
def update_groups(username):
    try:
        data = request.json
        groups = data.get('groups', [])

        users_data = load_users()

        if username not in users_data['users']:
            return jsonify({'error': 'User not found'}), 404

        users_data['users'][username]['groups'] = groups
        save_users(users_data)
        restart_authelia()

        return jsonify({'success': True})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/users/<username>/services', methods=['PUT'])
def update_services(username):
    """Update individual services for a user"""
    try:
        data = request.json
        services = data.get('services', [])

        users_data = load_users()

        if username not in users_data['users']:
            return jsonify({'error': 'User not found'}), 404

        users_data['users'][username]['services'] = services
        save_users(users_data)
        restart_authelia()

        return jsonify({'success': True})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/users/<username>/reset-password', methods=['POST'])
def reset_password(username):
    try:
        users_data = load_users()

        if username not in users_data['users']:
            return jsonify({'error': 'User not found'}), 404

        # Generate new password and hash it
        new_password = generate_password()
        password_hash = hash_password(new_password)

        # Update the password in the database
        users_data['users'][username]['password'] = password_hash
        save_users(users_data)
        restart_authelia()

        return jsonify({
            'success': True,
            'username': username,
            'password': new_password
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

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