import {
  EC2Client,
  RunInstancesCommand,
  TerminateInstancesCommand,
  DescribeInstancesCommand,
  CreateSecurityGroupCommand,
  DeleteSecurityGroupCommand,
  AuthorizeSecurityGroupIngressCommand,
  DescribeSecurityGroupsCommand,
  CreateKeyPairCommand,
  DeleteKeyPairCommand,
  CreateTagsCommand,
} from "@aws-sdk/client-ec2";
import { STSClient, GetCallerIdentityCommand } from "@aws-sdk/client-sts";
import { createLogger } from "../lib/logger";
import { decrypt } from "../lib/encryption";
import type { User } from "@prisma/client";

const log = createLogger("EC2");

const UBUNTU_22_04_AMI: Record<string, string> = {
  "us-east-1": "ami-0e86e20dae9224db8",
  "us-east-2": "ami-0a0d9cf81c479446a",
  "us-west-1": "ami-05c969369880fa2c2",
  "us-west-2": "ami-03f8acd418785369b",
  "eu-west-1": "ami-0694d931cee176e7d",
  "eu-west-2": "ami-0f3d9639a5674d559",
  "eu-west-3": "ami-022e307f4b9e39f45",
  "eu-central-1": "ami-0faab6bdbac9486fb",
  "ap-southeast-1": "ami-0823c236601fef765",
  "ap-southeast-2": "ami-07620139298af599e",
  "ap-northeast-1": "ami-0b7546e839d7ace12",
  "ap-northeast-2": "ami-042e76978adeb8c48",
  "ap-south-1": "ami-076e3a557efe1aa9c",
  "sa-east-1": "ami-0eed58016fbe42de3",
  "ca-central-1": "ami-024f768de9e73d4f4",
  "eu-north-1": "ami-00381a880aa48c6c6",
  "me-south-1": "ami-09574f34b8dcd2eac",
  "af-south-1": "ami-08fdcf06b39fe83ec",
};

export function makeEc2Client(user: User): EC2Client {
  const accessKeyId = user.awsAccessKeyId ? decrypt(user.awsAccessKeyId) : "";
  const secretAccessKey = user.awsSecretAccessKey ? decrypt(user.awsSecretAccessKey) : "";
  return new EC2Client({
    region: user.awsRegion!,
    credentials: { accessKeyId, secretAccessKey },
  });
}

export function makeStsClient(user: User): STSClient {
  const accessKeyId = user.awsAccessKeyId ? decrypt(user.awsAccessKeyId) : "";
  const secretAccessKey = user.awsSecretAccessKey ? decrypt(user.awsSecretAccessKey) : "";
  return new STSClient({
    region: user.awsRegion!,
    credentials: { accessKeyId, secretAccessKey },
  });
}

export async function validateAwsCredentials(user: User): Promise<{ success: boolean; arn?: string; error?: string }> {
  try {
    const accessKeyId = user.awsAccessKeyId ? decrypt(user.awsAccessKeyId) : "";
    const region = user.awsRegion ?? "";
    log.info("validateAwsCredentials: signer inputs", {
      accessKeyId: JSON.stringify(accessKeyId),
      accessKeyIdLen: accessKeyId.length,
      accessKeyIdCharCodes: [...accessKeyId].map((c) => c.charCodeAt(0)),
      region: JSON.stringify(region),
      regionCharCodes: [...region].map((c) => c.charCodeAt(0)),
    });
    const sts = makeStsClient(user);
    const res = await sts.send(new GetCallerIdentityCommand({}));
    return { success: true, arn: res.Arn };
  } catch (e: any) {
    return { success: false, error: e.message };
  }
}

export async function createKeyPairAndGetPrivateKey(ec2: EC2Client, keyName: string): Promise<{ privateKey: string }> {
  const res = await ec2.send(new CreateKeyPairCommand({
    KeyName: keyName,
    KeyType: "rsa",
    KeyFormat: "pem",
  }));
  return { privateKey: res.KeyMaterial! };
}

export async function createPreviewSecurityGroup(ec2: EC2Client, groupName: string, port: number): Promise<string> {
  const describe = await ec2.send(new DescribeSecurityGroupsCommand({
    Filters: [{ Name: "group-name", Values: [groupName] }],
  }));
  if (describe.SecurityGroups && describe.SecurityGroups.length > 0) {
    return describe.SecurityGroups[0].GroupId!;
  }

  const res = await ec2.send(new CreateSecurityGroupCommand({
    GroupName: groupName,
    Description: `PP Preview security group: ${groupName}`,
  }));
  const groupId = res.GroupId!;

  const ingress: any[] = [
    {
      IpProtocol: "tcp",
      FromPort: 22,
      ToPort: 22,
      IpRanges: [{ CidrIp: "0.0.0.0/0" }],
    },
  ];
  if (port !== 22) {
    ingress.push({
      IpProtocol: "tcp",
      FromPort: port,
      ToPort: port,
      IpRanges: [{ CidrIp: "0.0.0.0/0" }],
    });
  }

  await ec2.send(new AuthorizeSecurityGroupIngressCommand({
    GroupId: groupId,
    IpPermissions: ingress,
  }));
  return groupId;
}

const BOOTSTRAP_SCRIPT = `#!/bin/bash
set -euo pipefail
exec > >(tee -a /var/log/pp-bootstrap.log) 2>&1
trap 'touch /var/lib/pp-bootstrap-failed' ERR

# Ubuntu's apt-daily / unattended-upgrades timers fire on boot and hold the apt
# locks, which makes apt-get die immediately ("Could not get lock") — not just
# for our bootstrap, but for any user setup command that shells out to apt (e.g.
# the get.docker.com install script). "mask --now" stops the units if they are
# already running AND blocks them from being re-triggered for the life of this
# ephemeral instance. Then wait out anything mid-flight before each apt-get.
systemctl mask --now apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service unattended-upgrades.service 2>/dev/null || true
apt_wait() {
  for i in $(seq 1 150); do
    fuser /var/lib/dpkg/lock-frontend /var/lib/dpkg/lock /var/lib/apt/lists/lock /var/cache/apt/archives/lock >/dev/null 2>&1 && sleep 2 || break
  done
}

apt_wait
apt-get update -y
apt_wait
apt-get install -y ca-certificates curl git unzip

touch /var/lib/pp-bootstrap-done
`;

export async function launchInstance(opts: {
  ec2: EC2Client;
  region: string;
  instanceType: string;
  keyName: string;
  securityGroupId: string;
  tags: Record<string, string>;
}): Promise<string> {
  const ami = UBUNTU_22_04_AMI[opts.region] ?? UBUNTU_22_04_AMI["us-east-1"];
  const tagSpecs = Object.entries(opts.tags).map(([k, v]) => ({ Key: k, Value: v }));
  tagSpecs.push({ Key: "Name", Value: `pp-preview-${opts.tags["pp:previewId"]}` });

  const res = await opts.ec2.send(new RunInstancesCommand({
    ImageId: ami,
    InstanceType: opts.instanceType as any,
    MinCount: 1,
    MaxCount: 1,
    KeyName: opts.keyName,
    SecurityGroupIds: [opts.securityGroupId],
    UserData: Buffer.from(BOOTSTRAP_SCRIPT).toString("base64"),
    TagSpecifications: [
      { ResourceType: "instance", Tags: tagSpecs },
    ],
  }));

  return res.Instances![0].InstanceId!;
}

export async function waitForInstanceRunning(ec2: EC2Client, instanceId: string, maxWaitMs = 300_000): Promise<string> {
  const start = Date.now();
  while (Date.now() - start < maxWaitMs) {
    try {
      const res = await ec2.send(new DescribeInstancesCommand({
        InstanceIds: [instanceId],
      }));
      const inst = res.Reservations?.[0]?.Instances?.[0];
      if (inst?.State?.Name === "running" && inst.PublicIpAddress) {
        return inst.PublicIpAddress;
      }
    } catch (e: any) {
      // AWS eventual consistency: instance may not be visible immediately after RunInstances
      if (e.name !== "InvalidInstanceID.NotFound") throw e;
    }
    await sleep(5000);
  }
  throw new Error(`Instance ${instanceId} did not reach running state within timeout`);
}

export async function terminateInstance(ec2: EC2Client, instanceId: string): Promise<void> {
  await ec2.send(new TerminateInstancesCommand({ InstanceIds: [instanceId] }));
}

// Waits for an instance to reach the `terminated` state so its ENI is released
// and the security group no longer has a dependent object. Resolves (rather than
// throwing) on timeout so callers can still attempt SG deletion with retries.
export async function waitForInstanceTerminated(
  ec2: EC2Client,
  instanceId: string,
  timeoutMs = 180_000,
): Promise<void> {
  const deadline = Date.now() + timeoutMs;
  while (Date.now() < deadline) {
    try {
      const res = await ec2.send(new DescribeInstancesCommand({
        InstanceIds: [instanceId],
      }));
      const state = res.Reservations?.[0]?.Instances?.[0]?.State?.Name;
      if (!state || state === "terminated") return;
    } catch (e: any) {
      // Instance record aged out / never existed — nothing left to wait on.
      if (e.name === "InvalidInstanceID.NotFound") return;
      throw e;
    }
    await sleep(5000);
  }
}

export async function deleteKeyPairAws(ec2: EC2Client, keyName: string): Promise<void> {
  try {
    await ec2.send(new DeleteKeyPairCommand({ KeyName: keyName }));
  } catch (e) {
    log.warn({ e, keyName }, "Failed to delete key pair");
  }
}

export async function deleteSecurityGroupAws(ec2: EC2Client, groupName: string): Promise<void> {
  try {
    const describe = await ec2.send(new DescribeSecurityGroupsCommand({
      Filters: [{ Name: "group-name", Values: [groupName] }],
    }));
    const groupId = describe.SecurityGroups?.[0]?.GroupId;
    if (!groupId) return;

    // The terminated instance's ENI can linger for a while after the instance
    // itself is gone; DeleteSecurityGroup fails with DependencyViolation until
    // that ENI is released. Retry with backoff instead of failing outright.
    const maxAttempts = 12;
    for (let attempt = 1; attempt <= maxAttempts; attempt++) {
      try {
        await ec2.send(new DeleteSecurityGroupCommand({ GroupId: groupId }));
        return;
      } catch (e: any) {
        if (e.name === "InvalidGroup.NotFound") return;
        if (e.name === "DependencyViolation" && attempt < maxAttempts) {
          await sleep(10_000);
          continue;
        }
        throw e;
      }
    }
  } catch (e) {
    log.warn({ e, groupName }, "Failed to delete security group");
  }
}

export async function describeAllManagedInstances(ec2: EC2Client): Promise<any[]> {
  const res = await ec2.send(new DescribeInstancesCommand({
    Filters: [
      { Name: "tag:pp:managed", Values: ["true"] },
      { Name: "instance-state-name", Values: ["running", "pending", "stopping"] },
    ],
  }));
  return (res.Reservations ?? []).flatMap(r => r.Instances ?? []);
}

function sleep(ms: number) {
  return new Promise(r => setTimeout(r, ms));
}
