import datetime
import json
import os
import pickle
import random
import string
import tempfile
import time

import boto3
import pandas as pd
from awswrangler import s3


def generate_random_id(length=8):
    # Combine uppercase letters, lowercase letters, and digits
    characters = string.ascii_letters + string.digits
    # Use random.choices to generate a list of random characters of the specified length
    random_id = "".join(random.choices(characters, k=length))
    return random_id


class CustomJSONEncoder(json.JSONEncoder):
    """Custom encoder to convert datetime to str if needed."""

    def default(self, obj):
        if isinstance(obj, (pd.Timestamp, datetime.date)):
            return obj.isoformat()  # Convert to ISO 8601 string
        return super().default(obj)


def write_content_to_s3_or_local(content, output_path):

    extension = output_path.split(".")[-1]  # will be pkl or json

    if output_path.startswith("s3"):

        if extension == "json":
            with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file:
                json.dump(content, temp_file, indent=4)
                temp_file_path = temp_file.name
        elif extension == "pkl":
            with tempfile.NamedTemporaryFile(mode="wb", delete=False) as temp_file:
                pickle.dump(content, temp_file)
                temp_file_path = temp_file.name
        s3.upload(temp_file_path, output_path)
        os.remove(temp_file_path)

    else:
        if extension == "json":
            with open(output_path, "wb") as file:
                json.dump(content, file, indent=4)
        if extension == "pkl":
            with open(output_path, "wb") as file:
                pickle.dump(content, file)


def read_content_from_s3_or_local(path: str):
    extension = path.split(".")[-1]  # will be pkl or json
    read_from_s3 = path.startswith("s3")

    if read_from_s3:
        with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file:
            temp_file_path = temp_file.name
            s3.download(path, temp_file_path)
    else:
        temp_file_path = path

    with open(temp_file_path, "rb") as file:
        if extension == "json":
            data = json.load(file)
        elif extension == "pkl":
            data = pickle.load(file)

    if read_from_s3:
        os.remove(temp_file_path)
    return data


def start_ec2_instance(region_name: str, instance_id: str):
    ec2 = boto3.client("ec2", region_name=region_name)
    response = ec2.describe_instances(InstanceIds=[instance_id])
    state = response["Reservations"][0]["Instances"][0]["State"]["Name"]
    if state == "stopping":
        # Wait until the instance is stopped
        while True:
            response = ec2.describe_instances(InstanceIds=[instance_id])
            state = response["Reservations"][0]["Instances"][0]["State"]["Name"]
            if state == "stopped":
                ec2.start_instances(InstanceIds=[instance_id])
                break
            time.sleep(5)  # Wait for 10 seconds before checking again
    else:
        ec2.start_instances(InstanceIds=[instance_id])