import json
import os
import tarfile
from tempfile import TemporaryDirectory

from django.conf import settings as s
from django.core.management.base import BaseCommand
from django.db import transaction

from epdb.models import EnviFormer


class Command(BaseCommand):
    def add_arguments(self, parser):
        parser.add_argument(
            "model",
            type=str,
            help="Model UUID of the Model to Dump",
        )
        parser.add_argument("--output", type=str)

    def package_dict_and_folder(self, dict_data, folder_path, output_path):
        with TemporaryDirectory() as tmpdir:
            dict_filename = os.path.join(tmpdir, "data.json")

            with open(dict_filename, "w", encoding="utf-8") as f:
                json.dump(dict_data, f, indent=2)

            with tarfile.open(output_path, "w:gz") as tar:
                tar.add(dict_filename, arcname="data.json")
                tar.add(folder_path, arcname=os.path.basename(folder_path))

            os.remove(dict_filename)

    @transaction.atomic
    def handle(self, *args, **options):
        output = options["output"]

        if os.path.exists(output):
            raise ValueError(f"Output file {output} already exists")

        model = EnviFormer.objects.get(uuid=options["model"])

        data = {
            "uuid": str(model.uuid),
            "name": model.name,
            "description": model.description,
            "kv": model.kv,
            "data_packages_uuids": [str(p.uuid) for p in model.data_packages.all()],
            "eval_packages_uuids": [str(p.uuid) for p in model.data_packages.all()],
            "threshold": model.threshold,
            "eval_results": model.eval_results,
            "multigen_eval": model.multigen_eval,
            "model_status": model.model_status,
        }

        model_folder = os.path.join(s.MODEL_DIR, "enviformer", str(model.uuid))

        self.package_dict_and_folder(data, model_folder, output)