from __future__ import annotations
import os
import io
from uuid import uuid4
from typing import List, Tuple

from django.core.management.base import BaseCommand, CommandError
from django.core.files import File
from django.contrib.auth import get_user_model

from PIL import Image

from anatomy.models import AnatomyQuestion, Exam, ExamQuestionDetail, QuestionType, Structure
from generic.models import Modality as GenericModality  # type: ignore


class Command(BaseCommand):
    help = (
        "Create anatomy questions and assign them to an exam."
        "Requires --exam-id and --image-path (path readable inside the container)."
    )

    def add_arguments(self, parser):
        parser.add_argument("--exam-id", type=int, required=False, help="ID of the anatomy.Exam to assign questions to (if omitted a new exam will be created)")
        parser.add_argument("--exam-name", type=str, default=None, help="Name to use when creating a new exam if --exam-id omitted")
        parser.add_argument("--count", type=int, default=1, help="Number of questions to create")
        parser.add_argument("--image-path", type=str, required=False, help="Path to image file to use for each created question")
        parser.add_argument("--generate-image", action="store_true", help="Generate a placeholder image automatically if --image-path is not provided or not found")
        parser.add_argument("--image-size", type=str, default="800x600", help="Generated image size WIDTHxHEIGHT, e.g. 800x600")
        parser.add_argument("--image-color", type=str, default="#777777", help="Background color for generated image (hex) e.g. #ffffff")
        parser.add_argument("--description", type=str, default="Generated question {i}", help="Description template; use {i} for index")
        parser.add_argument("--question-type-id", type=int, default=None, help="Optional QuestionType id to assign")
        parser.add_argument("--modality-id", type=int, default=None, help="Optional Modality id to assign (generic.Modality id)")
        parser.add_argument("--structure-id", type=int, default=None, help="Optional Structure id to assign")
        parser.add_argument("--author-id", type=int, action="append", help="Author user id; repeatable")
        parser.add_argument("--start-sort", type=int, default=1, help="Starting sort_order for the first question")

    def handle(self, *args, **options):
        exam_id = options.get("exam_id")
        count = options["count"]
        image_path = options.get("image_path")
        generate_image = options.get("generate_image")
        image_size_str = options.get("image_size") or "800x600"
        image_color = options.get("image_color") or "#777777"
        description_template = options["description"]
        qtype_id = options.get("question_type_id")
        modality_id = options.get("modality_id")
        structure_id = options.get("structure_id")
        author_ids: List[int] = options.get("author_id") or []
        sort = options.get("start_sort") or 1

        exam = None
        if exam_id:
            try:
                exam = Exam.objects.get(pk=exam_id)
            except Exam.DoesNotExist:
                raise CommandError(f"Exam with id {exam_id} does not exist")
        else:
            # create new exam
            exam_name = options.get("exam_name") or f"Auto Anatomy Exam {uuid4().hex[:8]}"
            exam = Exam(name=exam_name, active=True, open_access=True)
            exam.save()
            self.stdout.write(self.style.SUCCESS(f"Created exam id={exam.id} name='{exam.name}'"))

        # validate image path if provided
        if image_path and not os.path.exists(image_path):
            if not generate_image:
                raise CommandError(f"Image file not found at path: {image_path}")
            else:
                image_path = None

        # Resolve optional foreign keys
        qtype = None
        if qtype_id:
            qtype = QuestionType.objects.filter(pk=qtype_id).first()
            if not qtype:
                raise CommandError(f"QuestionType id {qtype_id} not found")

        modality = None
        if modality_id:
            # The project stores Modality in generic.models.Modality
            try:
                from generic.models import Modality as GModality

                modality = GModality.objects.filter(pk=modality_id).first()
                if not modality:
                    raise CommandError(f"Modality id {modality_id} not found in generic.Modality")
            except Exception as e:
                raise CommandError(str(e))

        structure = None
        if structure_id:
            from atlas.models import Structure as AtlasStructure

            structure = AtlasStructure.objects.filter(pk=structure_id).first()
            if not structure:
                raise CommandError(f"Structure id {structure_id} not found")

        User = get_user_model()
        authors = []
        for uid in author_ids:
            u = User.objects.filter(pk=uid).first()
            if not u:
                raise CommandError(f"Author user id {uid} not found")
            authors.append(u)

        # parse image size
        try:
            w, h = image_size_str.lower().split("x")
            image_size: Tuple[int, int] = (int(w), int(h))
        except Exception:
            raise CommandError("Invalid --image-size; expected WIDTHxHEIGHT, e.g. 800x600")

        created = []

        for i in range(1, count + 1):
            desc = description_template.format(i=i)

            # determine image source: either provided path or generated image
            unique_name = f"{uuid4().hex}_image.jpg"

            if image_path:
                # use provided file
                with open(image_path, "rb") as f:
                    django_file = File(f, name=unique_name)
                    img_file_for_save = django_file
                    # create and save question
                    q = AnatomyQuestion(
                        question_type=qtype,
                        description=desc,
                        answer_help="",
                    )

                    if modality:
                        q.modality = modality

                    if structure:
                        q.structure = structure

                    q.image.save(unique_name, img_file_for_save, save=False)
                    q.save()

            else:
                # generate an image in memory
                buf = io.BytesIO()
                img = Image.new("RGB", image_size, image_color)
                img.save(buf, format="JPEG")
                buf.seek(0)
                django_file = File(buf, name=unique_name)

                q = AnatomyQuestion(
                    question_type=qtype,
                    description=desc,
                    answer_help="",
                )

                if modality:
                    q.modality = modality

                if structure:
                    q.structure = structure

                q.image.save(unique_name, django_file, save=False)
                q.save()

            # authors
            if authors:
                q.author.set(authors)

            # add to exam via the many-to-many through relationship
            try:
                exam.exam_questions.add(q, through_defaults={"sort_order": sort})
            except TypeError:
                ExamQuestionDetail.objects.create(exam=exam, anatomyquestion=q, sort_order=sort)

            self.stdout.write(self.style.SUCCESS(f"Created question id={q.id} and added to exam id={exam.id} with sort_order={sort}"))
            created.append(q)
            sort += 1

        self.stdout.write(self.style.SUCCESS(f"Finished: created {len(created)} question(s)"))
