#!/bin/bash
# Setup Authentik SSO for AI Tax Agent using Blueprint Import

set -euo pipefail

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Configuration
# Load environment variables
if [ -f "infra/compose/.env" ]; then
    source "infra/compose/.env"
fi

DOMAIN=${DOMAIN:-local}
AUTHENTIK_URL="https://auth.${DOMAIN}"
AUTHENTIK_API_URL="$AUTHENTIK_URL/api/v3"
ADMIN_EMAIL="admin@${DOMAIN}"
ADMIN_PASSWORD="${AUTHENTIK_ADMIN_PASSWORD:-admin123}"
BOOTSTRAP_FILE="infra/authentik/bootstrap.yaml"

echo -e "${BLUE}🔧 Setting up Authentik SSO for AI Tax Agent using Blueprint Import...${NC}"
echo

# Function to wait for Authentik to be ready
wait_for_authentik() {
    echo -e "${YELLOW}⏳ Waiting for Authentik to be ready...${NC}"
    local max_attempts=60
    local attempt=1
    local host
    host=$(echo "$AUTHENTIK_URL" | sed -E 's#^https?://([^/]+).*$#\1#')
    local resolve=(--resolve "${host}:443:127.0.0.1")

    while [ $attempt -le $max_attempts ]; do
        code_setup=$(curl -ks "${resolve[@]}" -o /dev/null -w '%{http_code}' "$AUTHENTIK_URL/if/flow/initial-setup/" || true)
        code_login=$(curl -ks "${resolve[@]}" -o /dev/null -w '%{http_code}' "$AUTHENTIK_URL/if/flow/default-authentication-flow/" || true)
        code_root=$(curl -ks "${resolve[@]}" -o /dev/null -w '%{http_code}' "$AUTHENTIK_URL/" || true)
        if [[ "$code_setup" == "404" ]]; then
            if [[ "$code_login" =~ ^(200|302|401)$ || "$code_root" =~ ^(200|302|401)$ ]]; then
                echo -e "${GREEN}✅ Authentik is ready!${NC}"; return 0
            fi
        fi
        if [[ "$code_setup" =~ ^(200|302|401)$ || "$code_login" =~ ^(200|302|401)$ || "$code_root" =~ ^(200|302|401)$ ]]; then
            echo -e "${GREEN}✅ Authentik is ready!${NC}"; return 0
        fi
        echo -n "."
        sleep 5
        ((attempt++))
    done

    echo -e "${RED}❌ Authentik failed to start within expected time${NC}"
    return 1
}

# Function to generate secrets
generate_secrets() {
    echo -e "${YELLOW}🔑 Generating secure secrets...${NC}"

    # Generate client secrets if not already set
    if [ -z "${AUTHENTIK_API_CLIENT_SECRET:-}" ]; then
        export AUTHENTIK_API_CLIENT_SECRET=$(openssl rand -base64 32 | tr -d '=+/')
        echo "Generated API client secret"
    fi

    if [ -z "${AUTHENTIK_GRAFANA_CLIENT_SECRET:-}" ]; then
        export AUTHENTIK_GRAFANA_CLIENT_SECRET=$(openssl rand -base64 32 | tr -d '=+/')
        echo "Generated Grafana client secret"
    fi

    if [ -z "${AUTHENTIK_SECRET_KEY:-}" ]; then
        export AUTHENTIK_SECRET_KEY=$(openssl rand -base64 50 | tr -d '=+/')
        echo "Generated Authentik secret key"
    fi

    echo -e "${GREEN}✅ Secrets generated${NC}"
}

# Function to get API token
get_api_token() {
    echo -e "${YELLOW}🔑 Getting API token...${NC}" >&2

    # Use bootstrap token if available and valid
    if [ -n "${AUTHENTIK_BOOTSTRAP_TOKEN:-}" ] && [ "$AUTHENTIK_BOOTSTRAP_TOKEN" != "ak-bootstrap-token" ]; then
        echo "$AUTHENTIK_BOOTSTRAP_TOKEN"
        return 0
    fi

    # Try to get token via API (requires manual setup first)
    local token_response
    token_response=$(curl -ks -X POST "$AUTHENTIK_API_URL/core/tokens/" \
        -H "Content-Type: application/json" \
        -u "$ADMIN_EMAIL:$ADMIN_PASSWORD" \
        -d '{
            "identifier": "ai-tax-agent-setup",
            "description": "Setup token for AI Tax Agent",
            "expires": "2025-12-31T23:59:59Z"
        }' 2>/dev/null || echo "")

    if [ -n "$token_response" ]; then
        echo "$token_response" | python3 -c "import sys, json; print(json.load(sys.stdin)['key'])" 2>/dev/null || echo ""
    else
        echo ""
    fi
}

# Function to import blueprint
import_blueprint() {
    local token="$1"

    echo -e "${YELLOW}📋 Importing Authentik blueprint...${NC}"

    if [ ! -f "$BOOTSTRAP_FILE" ]; then
        echo -e "${RED}❌ Bootstrap file not found: $BOOTSTRAP_FILE${NC}"
        return 1
    fi

    # Create blueprint instance
    local blueprint_response
    blueprint_response=$(curl -k -X POST "$AUTHENTIK_API_URL/managed/blueprints/" \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $token" \
        -d '{
            "name": "AI Tax Agent Bootstrap",
            "path": "ai-tax-agent-bootstrap.yaml",
            "context": {},
            "enabled": true
        }' 2>/dev/null || echo "")

    local blueprint_pk
    blueprint_pk=$(echo "$blueprint_response" | python3 -c "import sys, json; print(json.load(sys.stdin).get('pk', ''))" 2>/dev/null || echo "")

    if [ -z "$blueprint_pk" ]; then
        echo -e "${YELLOW}⚠️  Could not create blueprint. It might already exist. Trying to find it...${NC}"
        local existing_bp
        existing_bp=$(curl -k -X GET "$AUTHENTIK_API_URL/managed/blueprints/?name=AI%20Tax%20Agent%20Bootstrap" \
            -H "Authorization: Bearer $token" 2>/dev/null || echo "")

        blueprint_pk=$(echo "$existing_bp" | python3 -c "import sys, json; print(json.load(sys.stdin)['results'][0]['pk'])" 2>/dev/null || echo "")
    fi

    if [ -n "$blueprint_pk" ]; then
        echo -e "${GREEN}✅ Blueprint created with ID: $blueprint_pk${NC}"

        # Apply the blueprint
        echo -e "${YELLOW}🔄 Applying blueprint...${NC}"
        local apply_response
        apply_response=$(curl -k -X POST "$AUTHENTIK_API_URL/managed/blueprints/$blueprint_pk/apply/" \
            -H "Content-Type: application/json" \
            -H "Authorization: Bearer $token" \
            -d '{}' 2>/dev/null || echo "")

        echo -e "${GREEN}✅ Blueprint applied successfully${NC}"

        # Force-sync the Outpost token
        # The blueprint might fail to update the token for the existing embedded outpost, so we do it explicitly.
        echo -e "${YELLOW}🔄 Syncing Outpost token...${NC}"
        if docker exec -i apa-authentik-server python3 /manage.py shell -c "
from authentik.outposts.models import Outpost
from authentik.core.models import Token
import os

try:
    token_key = os.environ.get('AUTHENTIK_OUTPOST_TOKEN')
    if token_key:
        o = Outpost.objects.get(name='authentik Embedded Outpost')
        t = Token.objects.get(pk=o.token.pk)
        if t.key != token_key:
            t.key = token_key
            t.save()
            print('Token updated')
        else:
            print('Token already matches')
    else:
        print('No AUTHENTIK_OUTPOST_TOKEN found in environment')
except Exception as e:
    print(f'Error updating token: {e}')
    exit(1)
" > /dev/null; then
            echo -e "${GREEN}✅ Outpost token synced${NC}"
            # Restart outpost to pick up changes if needed (though it reads from env, so mostly for connection retry)
            docker restart apa-authentik-outpost > /dev/null 2>&1 || true
        else
            echo -e "${RED}❌ Failed to sync Outpost token${NC}"
        fi

    else
        echo -e "${RED}❌ Failed to create blueprint${NC}"
        return 1
    fi
}

# Function to check blueprint status
check_blueprint_status() {
    local token="$1"

    echo -e "${YELLOW}🔍 Checking blueprint status...${NC}"

    local blueprints_response
    blueprints_response=$(curl -s -X GET "$AUTHENTIK_API_URL/managed/blueprints/" \
        -H "Authorization: Bearer $token" 2>/dev/null || echo "")

    if [ -n "$blueprints_response" ]; then
        echo "$blueprints_response" | python3 -c "
import sys, json
try:
    data = json.load(sys.stdin)
    for bp in data.get('results', []):
        print(f\"Blueprint: {bp['name']} - Status: {'Enabled' if bp['enabled'] else 'Disabled'}\")
except:
    print('Could not parse blueprint status')
" 2>/dev/null || echo "Could not check blueprint status"
    fi
}



# Main setup function
main() {
    # Generate secrets first
    generate_secrets

    # Check if Authentik is running
    if ! wait_for_authentik; then
        echo -e "${RED}❌ Authentik is not accessible. Please ensure it's running.${NC}"
        exit 1
    fi

    # Check if initial setup is needed (only if we don't have a token)
    if [ -z "${AUTHENTIK_BOOTSTRAP_TOKEN:-}" ] || [ "$AUTHENTIK_BOOTSTRAP_TOKEN" == "ak-bootstrap-token" ]; then
        local host
        host=$(echo "$AUTHENTIK_URL" | sed -E 's#^https?://([^/]+).*$#\1#')
        local resolve=(--resolve "${host}:443:127.0.0.1")
        local setup_code
        setup_code=$(curl -ks "${resolve[@]}" -o /dev/null -w '%{http_code}' "$AUTHENTIK_URL/if/flow/initial-setup/" || true)

        if [[ "$setup_code" == "200" ]]; then
            echo -e "${YELLOW}📋 Initial Authentik setup required:${NC}"
            echo -e "  1. Open ${BLUE}https://auth.${DOMAIN}/if/flow/initial-setup/${NC}"
            echo -e "  2. Complete the setup wizard with admin user"
            echo -e "  3. Re-run this script after setup is complete"
            echo
            echo -e "${BLUE}💡 Tip: Use these credentials:${NC}"
            echo -e "  • Email: ${BLUE}$ADMIN_EMAIL${NC}"
            echo -e "  • Password: ${BLUE}$ADMIN_PASSWORD${NC}"
            return 0
        fi
    fi

    # Try to get API token
    local api_token
    api_token=$(get_api_token)

    if [ -n "$api_token" ]; then
        echo -e "${GREEN}🔑 API token obtained, proceeding with blueprint import...${NC}"

        # Import the blueprint configuration
        if import_blueprint "$api_token"; then
            echo -e "${GREEN}🎉 Authentik configuration imported successfully!${NC}"

            # Check status
            check_blueprint_status "$api_token"

            # Display client secrets for configuration
            echo
            echo -e "${BLUE}🔑 Client Secrets (save these for service configuration):${NC}"
            echo -e "  • API Client Secret: ${YELLOW}${AUTHENTIK_API_CLIENT_SECRET}${NC}"
            echo -e "  • Grafana Client Secret: ${YELLOW}${AUTHENTIK_GRAFANA_CLIENT_SECRET}${NC}"

        else
            echo -e "${RED}❌ Blueprint import failed${NC}"
            exit 1
        fi
    else
        echo -e "${YELLOW}📋 Could not obtain API token. Manual configuration required:${NC}"
        echo -e "  1. Open ${BLUE}https://auth.local.lan${NC} and log in as admin"
        echo -e "  2. Go to Admin Interface > Tokens"
        echo -e "  3. Create a new token and set AUTHENTIK_BOOTSTRAP_TOKEN in .env"
        echo -e "  4. Re-run this script"
    fi

    echo
    echo -e "${BLUE}🔗 Access URLs:${NC}"
    echo -e "  • Authentik Admin: ${BLUE}https://auth.local.lan${NC}"
    echo -e "  • API Gateway: ${BLUE}https://api.local.lan${NC}"
    echo -e "  • Grafana: ${BLUE}https://grafana.local.lan${NC}"
    echo -e "  • Review Portal: ${BLUE}https://review.local.lan${NC}"
}

# Run main function
main "$@"
