# myapp/management/commands/mycustomcommand.py
from django.core.management.base import BaseCommand
from tqdm import tqdm
from fwd_api.models import SubscriptionRequestFile, SubscriptionRequest
from fwd_api.utils.accuracy import predict_result_to_ready
import traceback
import copy

class Command(BaseCommand):
    help = 'Refactor database for image level'

    def add_arguments(self, parser):
        # Add your command-line arguments here
        parser.add_argument('test', type=str, help='Value for the argument')


    def process_request(self, request):
        if len(request.request_id.split(".")[0].split("_")) < 2:
            return
        images = SubscriptionRequestFile.objects.filter(request=request)
        time_cost = {"imei": [], "invoice": [], "all": []}
        if request.ai_inference_profile is None:
            time_cost["imei"] = [-1 for _ in range(len(images))]
            time_cost["invoice"] = [-1]
            time_cost["all"] = [-1]
        else:
            for k, v in request.ai_inference_profile.items():
                time_cost[k.split("_")[0]].append(v["inference"][1][0] - v["inference"][0] + (v["postprocess"][1]-v["postprocess"][0]))
        for i, image in enumerate(images):
            # temp_imei_SAP_20240127223644_a493434edbf84fc08aeb87ef6cdde102_0.jpg
            try:
                image.index_in_request = int(image.file_name.split(".")[0].split("_")[-1]) if len(image.file_name.split(".")[0].split("_")) > 4 else 0
                image.doc_type = image.file_name.split(".")[0].split("_")[1] if len(image.file_name.split(".")[0].split("_")) > 4 else "all"
                image.processing_time = time_cost[image.doc_type][image.index_in_request]
                if not request.predict_result:
                    raise KeyError(f"Key predict_result not found in {request.request_id}")
                if request.predict_result.get("status", 200) != 200:
                    raise AttributeError(f"Failed request: {request.request_id}")
                _predict_result = copy.deepcopy(predict_result_to_ready(request.predict_result))
                _feedback_result = copy.deepcopy(request.feedback_result)
                _reviewed_result = copy.deepcopy(request.reviewed_result)

                if image.doc_type == "invoice":
                    _predict_result["imei_number"] = []
                    if _feedback_result:
                        _feedback_result["imei_number"] = []
                    else:
                        None
                    if _reviewed_result:
                        _reviewed_result["imei_number"] = []
                    else:
                        None
                else:
                    _predict_result = {"retailername": None, "sold_to_party": None, "purchase_date": [], "imei_number": [_predict_result["imei_number"][image.index_in_request]]}
                    _feedback_result = {"retailername": None, "sold_to_party": None, "purchase_date": None, "imei_number": [_feedback_result["imei_number"][image.index_in_request]]} if _feedback_result else None
                    _reviewed_result = {"retailername": None, "sold_to_party": None, "purchase_date": None, "imei_number": [_reviewed_result["imei_number"][image.index_in_request]]} if _reviewed_result else None
                image.predict_result = _predict_result
                image.feedback_result = _feedback_result
                image.reviewed_result = _reviewed_result
                image.save()
            except Exception as e:
                self.stdout.write(self.style.ERROR(f"Request: {request.request_id} failed with {e}"))
                print(traceback.format_exc())
                continue

    def handle(self, *args, **options):
        test = options['test']
        subcription_iter = SubscriptionRequest.objects.all()
        for request in tqdm(subcription_iter.iterator()):
            self.process_request(request)
        self.stdout.write(self.style.SUCCESS('Sample Django management command executed successfully!'))