#!/bin/env zx
import fs from "fs"
import { exit } from "process";

const targetNumberOfCores = 32


const regions = JSON.parse(await $`aws ec2 describe-regions`).Regions.map(r => r.RegionName);

console.log("Available regions:", regions);


const instanceMap = JSON.parse(fs.readFileSync("instance-types.json", "utf-8"))
let pricings = JSON.parse(fs.readFileSync("spot-pricing.json", "utf-8"))

console.log("Spot pricing data:", pricings)

const burstableBaselines = {
    "t2.nano": 0.05,
    "t2.micro": 0.10,
    "t2.small": 0.20,
    "t2.medium": 0.20,
    "t2.large": 0.30,
    "t2.xlarge": 0.225,
    "t2.2xlarge": 0.17,
    "t3.nano": 0.05,
    "t3.micro": 0.10,
    "t3.small": 0.20,
    "t3.medium": 0.20,
    "t3.large": 0.30,
    "t3.xlarge": 0.40,
    "t3.2xlarge": 0.40,
    "t3a.nano": 0.05,
    "t3a.micro": 0.10,
    "t3a.small": 0.20,
    "t3a.medium": 0.20,
    "t3a.large": 0.30,
    "t3a.xlarge": 0.40,
    "t3a.2xlarge": 0.40,
    "t4g.nano": 0.05,
    "t4g.micro": 0.10,
    "t4g.small": 0.20,
    "t4g.medium": 0.20,
    "t4g.large": 0.30,
    "t4g.xlarge": 0.40,
    "t4g.2xlarge": 0.40
}

const performanceData = {
    "r7iz": 3077,
    "i7ie": 2989,
    "m7a": 2876,
    "c7a": 2870,
    "r7a": 2867,
    "c7i": 2856,
    "r7i": 2839,
    "c7i-flex": 2836,
    "m7i": 2831,
    "m7i-flex": 2761,
    "c6a": 2584,
    "m6a": 2584,
    "r6a": 2581,
    "m5zn": 2493,
    "m6in": 2288,
    "m6id": 2287,
    "r6i": 2286,
    "c6id": 2285,
    "r6idn": 2284,
    "m6i": 2284,
    "m6idn": 2283,
    "r6id": 2282,
    "i4i": 2280,
    "c6i": 2278,
    "c6in": 2276,
    "z1d": 2243,
    "r6in": 2243,
    "c5a": 2052,
    "c5ad": 2031,
    "c5n": 1895,
    "c5": 1894,
    "c5d": 1876,
    "t2": 1775,
    "r5n": 1735,
    "m5n": 1720,
    "r5dn": 1718,
    "m5dn": 1718,
    "m5": 1716,
    "r5": 1714,
    "r5b": 1712,
    "m5d": 1701,
    "r5d": 1679,
    "r4": 1574,
    "r3": 1572,
    "m4": 1542,
    "t3": 1517,
    "m5ad": 1405,
    "r5ad": 1404,
    "m5a": 1403,
    "r5a": 1394,
    "t3a": 1342
}


let sortedInstanceTypes = []

for (const p of pricings) {
    const i = instanceMap[p.InstanceType + "-" + p.AvailabilityZone.slice(0, p.AvailabilityZone.length - 1)]
    if (!i) continue
    if (i.VCpuInfo.DefaultVCpus > targetNumberOfCores) continue
    const sp = JSON.parse(JSON.stringify(i))
    sp.SpotPrice = p.SpotPrice
    sp.AvailabilityZone = p.AvailabilityZone
    let baseline = 1
    if (p.InstanceType in burstableBaselines) {
        baseline = burstableBaselines[p.InstanceType]
        continue
    }
    if (!p.AvailabilityZone.startsWith("eu-central")) {
        continue
    }
    const performance = (performanceData[p.InstanceType.split(".")[0]] ?? 1000) / 1000
    sp.SpotPricePerCore = p.SpotPrice / sp.VCpuInfo.DefaultVCpus / baseline / performance
    sortedInstanceTypes.push(sp)
}

sortedInstanceTypes = sortedInstanceTypes.sort((a, b) => a.SpotPricePerCore - b.SpotPricePerCore)

console.log("20 cheapest instance types:")
for (let i = 0; i < Math.min(20, sortedInstanceTypes.length); i++) {
    console.log(sortedInstanceTypes[i].InstanceType, sortedInstanceTypes[i].SpotPricePerCore, sortedInstanceTypes[i].SpotPrice, "in region", sortedInstanceTypes[i].AvailabilityZone)
}

const cheapestInstance = sortedInstanceTypes[0]
console.log("Cheapest instance type:", cheapestInstance)

const targetNumberOfInstances = Math.ceil(targetNumberOfCores / cheapestInstance.VCpuInfo.DefaultVCpus)
const targetRegion = cheapestInstance.AvailabilityZone.slice(0, cheapestInstance.AvailabilityZone.length - 1)

console.log(`Requesting ${targetNumberOfInstances} instances of type ${cheapestInstance.InstanceType} in region ${cheapestInstance.AvailabilityZone} at ${cheapestInstance.SpotPrice}$/h`)

//exit(0)

let vpc = JSON.parse(await $`aws ec2 describe-vpcs --filters Name=isDefault,Values=true --region ${targetRegion}`).Vpcs[0]

if (!vpc) {
    console.log("No default VPC found, creating one")
    await $`aws ec2 create-default-vpc --region ${targetRegion}`
    vpc = JSON.parse(await $`aws ec2 describe-vpcs --filters Name=isDefault,Values=true --region ${targetRegion}`).Vpcs[0]
}

if (!vpc.Ipv6CidrBlockAssociationSet) {
    console.log("No IPv6 CIDR block found, associating one")
    await $`aws ec2 associate-vpc-cidr-block --vpc-id ${vpc.VpcId} --ipv6-cidr-block-network-border-group ${targetRegion} --amazon-provided-ipv6-cidr-block --region ${targetRegion}`
    vpc = JSON.parse(await $`aws ec2 describe-vpcs --filters Name=isDefault,Values=true --region ${targetRegion}`).Vpcs[0]
}

let ipv6CidrBlock = vpc.Ipv6CidrBlockAssociationSet.filter(c => c.Ipv6CidrBlockState.State == "associated")[0].Ipv6CidrBlock

let subnets = JSON.parse(await $`aws ec2 describe-subnets --filters Name=vpc-id,Values=${vpc.VpcId} --region ${targetRegion}`).Subnets

let targetRegionSubnetId = subnets.filter(s => s.AvailabilityZone == cheapestInstance.AvailabilityZone)[0].SubnetId

for (const subnet of subnets) {
    if (subnet.Ipv6CidrBlockAssociationSet.length == 0) {
        const parts = ipv6CidrBlock.split("/")[0].split(":")
        let prefix = parts.slice(0, 3).join(":") + ":"
        let lastPart = parts[3].slice(0, 2)
        while (lastPart.length < 3) {
            lastPart = lastPart + "0"
        }
        prefix += lastPart
        let i = 0
        while (subnets.some(s => s.Ipv6CidrBlockAssociationSet.some(c => c.Ipv6CidrBlock.startsWith(`${prefix}${i}`)))) {
            i++
        }
        const cidrBlock = `${prefix}${i}::/64`


        console.log(`No IPv6 CIDR block found for subnet ${subnet.SubnetId}, associating one`)
        await $`aws ec2 associate-subnet-cidr-block --subnet-id ${subnet.SubnetId} --ipv6-cidr-block ${cidrBlock} --region ${targetRegion}`
        subnet.Ipv6CidrBlockAssociationSet = JSON.parse(await $`aws ec2 describe-subnets --filters Name=subnet-id,Values=${subnet.SubnetId} --region ${targetRegion}`).Subnets[0].Ipv6CidrBlockAssociationSet
    }
}

let securityGroup = JSON.parse(await $`aws ec2 describe-security-groups --filters Name=group-name,Values=default --region ${targetRegion}`).SecurityGroups[0]

if (!securityGroup.IpPermissions.some(p => p.UserIdGroupPairs.length == 0)) {
    console.log("Allowing all inbound traffic to default security group")
    if (securityGroup.IpPermissions.length > 0) {
        for (const p of securityGroup.IpPermissions.filter(p => p.UserIdGroupPairs.length > 0)) {
            await $`aws ec2 revoke-security-group-ingress --group-id ${securityGroup.GroupId} --ip-permissions IpProtocol=-1,UserIdGroupPairs=[{GroupId=${p.UserIdGroupPairs[0].GroupId}}] --region ${targetRegion}`
        }
    }
    await $`aws ec2 authorize-security-group-ingress --group-id ${securityGroup.GroupId} --ip-permissions IpProtocol=-1,FromPort=-1,ToPort=-1,Ipv6Ranges=[{CidrIpv6=::/0}] --region ${targetRegion}`
}

let internetGateways = JSON.parse(await $`aws ec2 describe-internet-gateways --region ${targetRegion}`).InternetGateways

if (internetGateways.length == 0) {
    console.log("No internet gateway found, creating one")
    await $`aws ec2 create-internet-gateway --region ${targetRegion}`
    internetGateways = JSON.parse(await $`aws ec2 describe-internet-gateways --region ${targetRegion}`).InternetGateways
}
if (!internetGateways.some(g => g.Attachments.some(a => a.VpcId == vpc.VpcId))) {
    console.log("Attaching internet gateway to VPC")
    await $`aws ec2 attach-internet-gateway --internet-gateway-id ${internetGateways[0].InternetGatewayId} --vpc-id ${vpc.VpcId} --region ${targetRegion}`
    internetGateways = JSON.parse(await $`aws ec2 describe-internet-gateways --region ${targetRegion}`).InternetGateways
}

let internetGateway = internetGateways.filter(g => g.Attachments.some(a => a.VpcId == vpc.VpcId))[0]

let routeTable = JSON.parse(await $`aws ec2 describe-route-tables --filters Name=vpc-id,Values=${vpc.VpcId} --region ${targetRegion}`).RouteTables[0]

if (routeTable.Routes.some(r => r.DestinationIpv6CidrBlock == "::/0" && r.GatewayId !== internetGateway.InternetGatewayId)) {
    console.log("Removing incorrect route to internet gateway")
    await $`aws ec2 delete-route --route-table-id ${routeTable.RouteTableId} --destination-ipv6-cidr-block ::/0 --region ${targetRegion}`
    routeTable = JSON.parse(await $`aws ec2 describe-route-tables --filters Name=vpc-id,Values=${vpc.VpcId} --region ${targetRegion}`).RouteTables[0]
}

if (!routeTable.Routes.some(r => r.DestinationIpv6CidrBlock == "::/0")) {
    console.log("Adding route to internet gateway")
    await $`aws ec2 create-route --route-table-id ${routeTable.RouteTableId} --destination-ipv6-cidr-block ::/0 --gateway-id ${internetGateway.InternetGatewayId} --region ${targetRegion}`
    routeTable = JSON.parse(await $`aws ec2 describe-route-tables --filters Name=vpc-id,Values=${vpc.VpcId} --region ${targetRegion}`).RouteTables[0]
}

const publiKey = "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCrJ2vQmLUH9Q5aJJHTAT3thfw2u/6GMukFf5m7PCGfTCvqjnO+Qjex/p2KmeBpPDoLzlhoE4ROLTT96A1bRg2S41MnrZEd+gTUpnpVjNVWbnZHvEpQjKAKYDLGOxfhEM9webeXwZd37rmodg6p5EYPCsphXK2IziKxI9WS1CekgsLxtW7JE3HrRzSSt7/Wzo1TBOzoGD/KoW4S9AozQiD0ZEgPwj9n7yjyjtlORKAsaJQwylpUYKOGvSNqRQykCsKrBnigOURmPDzwI65jzZj1SyYsBtcCrgfLvqFmJGUD0GdN4A/64Q8gF+xbh7HRI1SDzbXowbeHoYLSjQqQEa9Z peeters@enterprise.fritz.box"

let keyPairs = JSON.parse(await $`aws ec2 describe-key-pairs --filters Name=key-name,Values=enterprise --region ${targetRegion}`).KeyPairs
if (keyPairs.length == 0) {
    console.log("No key pair found, importing...")
    await $`aws ec2 import-key-pair --key-name enterprise --public-key-material ${publiKey} --region ${targetRegion}`
}

const image = JSON.parse(await $`aws ec2 describe-images --owners amazon --filters "Name=name,Values=ubuntu-minimal/images/hvm-ssd-gp3/ubuntu-oracular-24.10-amd64-minimal-20250321" --region ${targetRegion}`).Images[0]

const dateFrom = new Date(Math.ceil(new Date().getTime() / 1000) * 1000)
const dateUntil = new Date(Math.ceil(new Date(dateFrom.getTime() + 10 * 60 * 60 * 1000).getTime() / 1000) * 1000)

const setupScript = fs.readFileSync("setup.sh", "utf8")
const setupScriptBase64 = Buffer.from(setupScript).toString("base64")

/**
 * @type {import("@aws-sdk/client-ec2").CreateFleetRequest}
 */
const fleet = {
    TargetCapacitySpecification: {
        TotalTargetCapacity: targetNumberOfInstances,
        DefaultTargetCapacityType: "spot"
    },
    IamFleetRole: "arn:aws:iam::228068582184:role/aws-ec2-spot-fleet-tagging-role",
    SpotOptions: {
        AllocationStrategy: "priceCapacityOptimized",
        MaxTotalPrice: "0.2",
    },
    TargetCapacity: targetNumberOfInstances,
    ValidFrom: dateFrom.toISOString(),
    ValidUntil: dateUntil.toISOString(),
    Type: "request",
    TerminateInstancesWithExpiration: true,

    LaunchTemplateConfigs: subnets.map(s => {
        /**
         * @type {import("@aws-sdk/client-ec2").FleetLaunchTemplateConfigRequest}
         */
        const template = {
            LaunchTemplateSpecification: {
                ImageId: image.ImageId,
                BlockDeviceMappings: [
                    {
                        DeviceName: "/dev/sda1",
                        Ebs: {
                            DeleteOnTermination: true,
                            SnapshotId: image.BlockDeviceMappings.filter(b => !!b.Ebs)[0].Ebs.SnapshotId,
                            VolumeSize: 8,
                            VolumeType: "gp3",
                            Encrypted: false
                        }
                    }
                ],
                KeyName: "enterprise",
                UserData: setupScriptBase64,
                NetworkInterfaces: [
                    {
                        DeviceIndex: 0,
                        SubnetId: s.SubnetId,
                        DeleteOnTermination: true,
                        Groups: [],
                        AssociatePublicIpAddress: false,
                        Ipv6AddressCount: 1,
                    }
                ],
                InstanceType: cheapestInstance.InstanceType,
                Placement: {
                    AvailabilityZone: s.AvailabilityZone,
                    Tenancy: "default",
                },
            },
        }
        return template
    }),
}

console.log("Fleet configuration:", fleet)
//exit(0)

console.log("Requesting spot fleet...")
await $`aws ec2 request-spot-fleet --spot-fleet-request-config ${JSON.stringify(fleet)} --region ${targetRegion}`