"""Launch the FelaTab x TabArena benchmark on EC2 (replaces the unreliable HF Jobs path).

Pushes the current benchmark scripts to the HF model repo, then boots an on-demand
c7i.8xlarge (32 vCPU / 64 GiB, ~$1.43/hr in us-west-2) whose user-data runs
ec2_run.sh: lite shakedown -> (guard) -> full suite -> upload artifacts to
lowdown-labs/fela-tab-tabarena-results -> poweroff (shutdown behavior = terminate).

Prereqs: `aws login` (SSO session), HF_TOKEN in env (write token for lowdown-labs).

Usage:
  HF_TOKEN=hf_... python launch_ec2.py                      # lite shakedown then full
  HF_TOKEN=hf_... python launch_ec2.py --phases lite        # shakedown only
  HF_TOKEN=hf_... python launch_ec2.py --phases full --datasets "houses,diamonds,..."  # shard

Note: HF_TOKEN is embedded in EC2 user-data (visible to anyone with
DescribeInstanceAttribute permission on this account). Rotate the token after the run
if that matters to you.
"""

import argparse
import os
import time
import urllib.request

import boto3
from huggingface_hub import HfApi, get_token

MODEL_REPO = "lowdown-labs/fela-tab"
INSTANCE_TYPE = "c7i.8xlarge"
KEY_NAME = "felatab-tabarena"
SG_NAME = "felatab-tabarena-ssh"
AMI_PARAM = "/aws/service/canonical/ubuntu/server/24.04/stable/current/amd64/hvm/ebs-gp3/ami-id"

USER_DATA_TMPL = """#!/bin/bash
set -euxo pipefail
export HOME=/root
export DEBIAN_FRONTEND=noninteractive
apt-get update -qq
apt-get install -y -qq git curl python3 python3-venv
curl -LsSf https://astral.sh/uv/install.sh | sh
export PATH="$HOME/.local/bin:$PATH"
uv venv --seed /work-bootvenv
uv pip install --python /work-bootvenv/bin/python -q huggingface_hub
mkdir -p /work
export HF_TOKEN="{hf_token}"
/work-bootvenv/bin/python -c "from huggingface_hub import snapshot_download; snapshot_download('{model_repo}', allow_patterns=['benchmark/tabarena/ec2_run.sh'], local_dir='/work/fela-tab')"
cat > /etc/felatab.env <<EOF
HF_TOKEN={hf_token}
PHASES={phases}
RUN_NAME={run_name}
DATASETS={datasets}
WITH_LGBM={with_lgbm}
SELF_TERMINATE=1
CHECKPOINT={checkpoint}
RESTORE={restore}
EOF
chmod 600 /etc/felatab.env
cat > /etc/systemd/system/felatab-benchmark.service <<'EOF'
[Unit]
Description=FelaTab TabArena benchmark
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
EnvironmentFile=/etc/felatab.env
ExecStart=/bin/bash /work/fela-tab/benchmark/tabarena/ec2_run.sh
User=root
StandardOutput=append:/work/logs_boot.log
StandardError=append:/work/logs_boot.log

[Install]
WantedBy=multi-user.target
EOF
systemctl daemon-reload
systemctl enable --now felatab-benchmark.service
"""


def ensure_key_pair(ec2, key_path):
    names = {k["KeyName"] for k in ec2.describe_key_pairs()["KeyPairs"]}
    if KEY_NAME in names:
        print(f"key pair {KEY_NAME} exists")
        return
    kp = ec2.create_key_pair(KeyName=KEY_NAME)
    with open(key_path, "w") as f:
        f.write(kp["KeyMaterial"])
    os.chmod(key_path, 0o400)
    print(f"created key pair {KEY_NAME} -> {key_path}")


def ensure_security_group(ec2):
    vpcs = ec2.describe_vpcs(Filters=[{"Name": "isDefault", "Values": ["true"]}])["Vpcs"]
    vpc_id = vpcs[0]["VpcId"]
    for sg in ec2.describe_security_groups(
        Filters=[{"Name": "group-name", "Values": [SG_NAME]}, {"Name": "vpc-id", "Values": [vpc_id]}]
    )["SecurityGroups"]:
        print(f"security group {SG_NAME} exists ({sg['GroupId']})")
        return sg["GroupId"]
    my_ip = urllib.request.urlopen("https://checkip.amazonaws.com").read().decode().strip()
    sg = ec2.create_security_group(GroupName=SG_NAME, Description="FelaTab TabArena runner (ssh)", VpcId=vpc_id)
    ec2.authorize_security_group_ingress(
        GroupId=sg["GroupId"],
        IpPermissions=[{"IpProtocol": "tcp", "FromPort": 22, "ToPort": 22,
                        "IpRanges": [{"CidrIp": f"{my_ip}/32", "Description": "launcher ip"}]}],
    )
    print(f"created security group {SG_NAME} ({sg['GroupId']}), ssh from {my_ip}/32")
    return sg["GroupId"]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--phases", default="lite,full", help="comma list: lite,full or just lite / full")
    ap.add_argument("--datasets", default="", help="optional csv shard (full mode only)")
    ap.add_argument("--run-name", default="felatab_ec2")
    ap.add_argument("--instance-type", default=INSTANCE_TYPE)
    ap.add_argument("--with-lgbm", default="0")
    ap.add_argument("--spot", action="store_true",
                    help="persistent spot with stop-on-interruption: EBS cache survives, "
                         "systemd service resumes the run on restart")
    ap.add_argument("--checkpoint", default="0", choices=["0", "1"],
                    help="upload results cache to the HF repo every 30 min")
    ap.add_argument("--restore", default="0", choices=["0", "1"],
                    help="download existing raw/<run-name> cache before running")
    args = ap.parse_args()

    token = os.environ.get("HF_TOKEN") or get_token()
    assert token, "need HF_TOKEN in env or a logged-in `hf auth login` session"
    api = HfApi(token=token)

    here = os.path.dirname(os.path.abspath(__file__))
    for fname in ("ec2_run.sh", "run_tabarena.py", "fela_ag_model.py"):
        api.upload_file(
            path_or_fileobj=os.path.join(here, fname),
            path_in_repo=f"benchmark/tabarena/{fname}",
            repo_id=MODEL_REPO, repo_type="model",
            commit_message=f"Update TabArena EC2 runner scripts ({fname})",
        )
        print(f"uploaded {fname} -> {MODEL_REPO}")

    ec2 = boto3.client("ec2")
    ssm = boto3.client("ssm")
    ami = ssm.get_parameter(Name=AMI_PARAM)["Parameter"]["Value"]
    print(f"AMI: {ami}")

    key_path = os.path.expanduser(f"~/.aws/{KEY_NAME}.pem")
    ensure_key_pair(ec2, key_path)
    sg_id = ensure_security_group(ec2)

    user_data = USER_DATA_TMPL.format(
        hf_token=token, model_repo=MODEL_REPO, phases=args.phases,
        run_name=args.run_name, datasets=args.datasets, with_lgbm=args.with_lgbm,
        checkpoint=args.checkpoint, restore=args.restore,
    )
    kwargs = dict(
        ImageId=ami,
        InstanceType=args.instance_type,
        KeyName=KEY_NAME,
        MinCount=1, MaxCount=1,
        NetworkInterfaces=[{"DeviceIndex": 0, "Groups": [sg_id], "AssociatePublicIpAddress": True}],
        BlockDeviceMappings=[{"DeviceName": "/dev/sda1",
                              "Ebs": {"VolumeSize": 200, "VolumeType": "gp3", "DeleteOnTermination": True}}],
        TagSpecifications=[{"ResourceType": "instance",
                            "Tags": [{"Key": "Name", "Value": f"felatab-tabarena-{args.run_name}"}]}],
        UserData=user_data,
    )
    if args.spot:
        # persistent spot, stop (not terminate) on interruption: the EBS results cache
        # survives and the systemd service resumes the run when capacity returns
        kwargs["InstanceMarketOptions"] = {
            "MarketType": "spot",
            "SpotOptions": {"SpotInstanceType": "persistent", "InstanceInterruptionBehavior": "stop"},
        }
        kwargs["InstanceInitiatedShutdownBehavior"] = "stop"
    else:
        kwargs["InstanceInitiatedShutdownBehavior"] = "terminate"
    resp = ec2.run_instances(**kwargs)
    iid = resp["Instances"][0]["InstanceId"]
    print(f"instance: {iid} ({args.instance_type}, phases={args.phases}, run={args.run_name})")

    ec2.get_waiter("instance_running").wait(InstanceIds=[iid])
    desc = ec2.describe_instances(InstanceIds=[iid])["Reservations"][0]["Instances"][0]
    ip = desc.get("PublicIpAddress", "<no public ip>")
    print(f"running at {ip}")
    print(f"watch:  ssh -i {key_path} ubuntu@{ip} 'tail -f /work/logs_boot.log'")
    print(f"status: aws ec2 describe-instances --instance-ids {iid} --query 'Reservations[0].Instances[0].State.Name'")
    print("artifacts land in https://huggingface.co/datasets/lowdown-labs/fela-tab-tabarena-results")


if __name__ == "__main__":
    main()