// Name: EC2 Security Group Modifier
// Description: Select an EC2 instance, view its security groups and inbound rules, and add your current IP to a chosen rule for easy SSH access.
// Author: anjaneyasivan
// GitHub: anjaneyasivan

import "@johnlindquist/kit"

const AWS_ACCESS_KEY_ID = await env("AWS_ACCESS_KEY_ID", {
  secret: true,
  placeholder: "Enter your AWS Access Key ID",
  hint: "Stored securely in ~/.kenv/.env",
})

const AWS_SECRET_ACCESS_KEY = await env("AWS_SECRET_ACCESS_KEY", {
  secret: true,
  placeholder: "Enter your AWS Secret Access Key",
  hint: "Stored securely in ~/.kenv/.env",
})

const AWS_REGION = await env(
  "AWS_REGION",
  async () =>
    await arg("Select your AWS region", [
      "us-east-1",
      "us-east-2",
      "us-west-1",
      "us-west-2",
      "eu-central-1",
      "eu-west-1",
      "eu-west-2",
      "eu-west-3",
      "ap-south-1",
      "ap-southeast-1",
      "ap-southeast-2",
      "ap-northeast-1",
      "ap-northeast-2",
      "sa-east-1",
      "ca-central-1",
    ])
)

const ec2Module =
  (await silentAttemptImport("@aws-sdk/client-ec2").catch(async () => {
    await npmInstall("@aws-sdk/client-ec2")
    return await attemptImport("@aws-sdk/client-ec2")
  })) || (await attemptImport("@aws-sdk/client-ec2"))

const {
  EC2Client,
  DescribeInstancesCommand,
  DescribeSecurityGroupsCommand,
  AuthorizeSecurityGroupIngressCommand,
} = ec2Module

const ec2 = new EC2Client({
  region: AWS_REGION,
  credentials: {
    accessKeyId: AWS_ACCESS_KEY_ID,
    secretAccessKey: AWS_SECRET_ACCESS_KEY,
  },
})

type InstanceSummary = {
  id: string
  name: string
  state: string
  publicIp?: string
  privateIp?: string
  sgList: { GroupId: string; GroupName?: string }[]
  raw: any
}

async function listInstances(): Promise<InstanceSummary[]> {
  let nextToken: string | undefined = undefined
  const result: InstanceSummary[] = []

  do {
    const res = await ec2.send(
      new DescribeInstancesCommand({
        NextToken: nextToken,
      })
    )
    nextToken = res.NextToken

    for (const r of res.Reservations || []) {
      for (const i of r.Instances || []) {
        const nameTag = (i.Tags || []).find(t => t.Key === "Name")?.Value || ""
        result.push({
          id: i.InstanceId || "",
          name: nameTag || i.InstanceId || "",
          state: i.State?.Name || "",
          publicIp: i.PublicIpAddress,
          privateIp: i.PrivateIpAddress,
          sgList: (i.SecurityGroups || []).map(sg => ({
            GroupId: sg.GroupId || "",
            GroupName: sg.GroupName,
          })),
          raw: i,
        })
      }
    }
  } while (nextToken)

  return result.sort((a, b) => a.name.localeCompare(b.name))
}

function permissionToLabel(p: any): string {
  const proto = p.IpProtocol === "-1" ? "all" : p.IpProtocol
  const port =
    p.IpProtocol === "-1"
      ? ""
      : typeof p.FromPort === "number" && typeof p.ToPort === "number"
      ? p.FromPort === p.ToPort
        ? `:${p.FromPort}`
        : `:${p.FromPort}-${p.ToPort}`
      : ""
  const ipv4 = (p.IpRanges || []).map((r: any) => r.CidrIp).filter(Boolean)
  const ipv6 = (p.Ipv6Ranges || []).map((r: any) => r.CidrIpv6).filter(Boolean)
  const ipsPreview = [...ipv4, ...ipv6].slice(0, 3).join(", ")
  const ipsCount = (ipv4.length + ipv6.length) || 0
  const counts = ipsCount ? ` (${ipsCount} ranges${ipsPreview ? `: ${ipsPreview}` : ""})` : ""
  return `${proto}${port}${counts}`
}

function permissionToIngressInput(p: any, cidrIp: string, description?: string) {
  const input: any = {
    IpProtocol: p.IpProtocol,
    IpRanges: [{ CidrIp: cidrIp, Description: description || "Added by Script Kit" }],
  }
  if (typeof p.FromPort === "number") input.FromPort = p.FromPort
  if (typeof p.ToPort === "number") input.ToPort = p.ToPort
  return input
}

async function getMyIp(): Promise<string> {
  try {
    const res = await get("https://api.ipify.org?format=json")
    return res.data?.ip
  } catch {
    const alt = await get("https://checkip.amazonaws.com")
    return String(alt.data || "").trim()
  }
}

const myIp = await getMyIp()
const myCidr = `${myIp}/32`

const instances = await listInstances()
if (!instances.length) {
  await div(md(`# No EC2 Instances Found in ${AWS_REGION}`))
  exit()
}

const instance = await arg<InstanceSummary>(
  {
    placeholder: "Select an EC2 Instance",
    enter: "Select",
    preview: async (input, state) => {
      const f = state?.focused?.value as InstanceSummary
      if (!f) return ""
      const sgs = f.sgList.map(s => `- ${s.GroupName || s.GroupId} (${s.GroupId})`).join("\n") || "None"
      return md(
        `# ${f.name}
- Instance ID: ${f.id}
- State: ${f.state}
- Public IP: ${f.publicIp || "—"}
- Private IP: ${f.privateIp || "—"}

## Attached Security Groups
${sgs}

> Your IP: ${myCidr}`
      )
    },
  },
  instances.map(i => ({
    name: `${i.name}  —  ${i.state}  —  ${i.publicIp || "no public ip"}`,
    description: `ID: ${i.id}`,
    value: i,
  }))
)

if (!instance.sgList.length) {
  await div(md(`# No Security Groups attached to instance ${instance.name}`))
  exit()
}

const sgChoice = await arg<{ GroupId: string; GroupName?: string }>(
  {
    placeholder: `Select a Security Group on ${instance.name}`,
    enter: "Select",
    preview: async (input, state) => {
      const sg = state?.focused?.value
      if (!sg) return ""
      return md(`## ${sg.GroupName || sg.GroupId}\nID: ${sg.GroupId}`)
    },
  },
  instance.sgList.map(sg => ({
    name: `${sg.GroupName || sg.GroupId}`,
    description: `ID: ${sg.GroupId}`,
    value: sg,
  }))
)

const sgDetails = await ec2.send(
  new DescribeSecurityGroupsCommand({
    GroupIds: [sgChoice.GroupId],
  })
)

const group = (sgDetails.SecurityGroups || [])[0]
if (!group) {
  await div(md(`# Unable to load Security Group ${sgChoice.GroupId}`))
  exit()
}

const inbound = group.IpPermissions || []
if (!inbound.length) {
  await div(md(`# No inbound rules on ${group.GroupName || group.GroupId}`))
  exit()
}

const permChoice = await arg<any>(
  {
    placeholder: `Select an inbound rule to add ${myCidr}`,
    enter: "Add my IP",
    preview: async (input, state) => {
      const p = state?.focused?.value
      if (!p) return ""
      const proto = p.IpProtocol === "-1" ? "all" : p.IpProtocol
      const port =
        p.IpProtocol === "-1"
          ? "all"
          : typeof p.FromPort === "number" && typeof p.ToPort === "number"
          ? p.FromPort === p.ToPort
            ? `${p.FromPort}`
            : `${p.FromPort}-${p.ToPort}`
          : "—"
      const ipv4 = (p.IpRanges || []).map((r: any) => `${r.CidrIp}${r.Description ? ` (${r.Description})` : ""}`).join("\n") || "None"
      const ipv6 = (p.Ipv6Ranges || []).map((r: any) => `${r.CidrIpv6}${r.Description ? ` (${r.Description})` : ""}`).join("\n") || "None"
      return md(
        `# Inbound Rule
- Protocol: ${proto}
- Ports: ${port}

## IPv4 Ranges
${ipv4}

## IPv6 Ranges
${ipv6}

> Will add: ${myCidr}`
      )
    },
  },
  inbound.map((p, idx) => ({
    name: permissionToLabel(p),
    description:
      p.IpProtocol === "-1"
        ? "All protocols"
        : typeof p.FromPort === "number" && typeof p.ToPort === "number"
        ? p.FromPort === p.ToPort
          ? `Port ${p.FromPort}`
          : `Ports ${p.FromPort}-${p.ToPort}`
        : "Ports —",
    value: p,
  }))
)

const alreadyHas =
  (permChoice.IpRanges || []).some((r: any) => r?.CidrIp === myCidr) ||
  (permChoice.Ipv6Ranges || []).some((r: any) => r?.CidrIpv6 === myCidr)

if (alreadyHas) {
  await notify(`IP ${myCidr} already present on selected rule`)
  await div(md(`✅ Your IP ${myCidr} is already present on the selected rule.`))
  exit()
}

try {
  const ingress = permissionToIngressInput(permChoice, myCidr, "Added by Script Kit for SSH access")
  await ec2.send(
    new AuthorizeSecurityGroupIngressCommand({
      GroupId: group.GroupId,
      IpPermissions: [ingress],
    })
  )

  await notify(`Added ${myCidr} to ${group.GroupName || group.GroupId}`)
  await div(
    md(
      `# Success
- Instance: ${instance.name} (${instance.id})
- Security Group: ${group.GroupName || group.GroupId}
- Added: ${myCidr}
- Rule: ${permissionToLabel(permChoice)}`
    )
  )
} catch (err: any) {
  const msg = err?.message || String(err)
  await div(
    md(
      `# Error Adding Ingress
${"```"}
${msg}
${"```"}`
    )
  )
  exit(1)
}