import io
import mimetypes
import os
import uuid
import zipfile
from datetime import datetime
from tempfile import NamedTemporaryFile

from django.conf import settings
from django.db.models import Max
from PIL import Image, UnidentifiedImageError

from apps.common.minio_client import minio_storage
from apps.common.utils.annotation_store import split_dataset_sample_annotation_content
from apps.core.models import AiDatasetSample

try:
    import cv2
except Exception:
    cv2 = None


DATASET_BUCKET = "dataset"
TEMPFILE_BUCKET = "tempfile"
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".webp", ".tif", ".tiff"}
VIDEO_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".wmv", ".flv", ".mpeg", ".mpg", ".webm"}
AUDIO_EXTENSIONS = {".mp3", ".wav", ".aac", ".m4a", ".flac", ".ogg"}
TEXT_EXTENSIONS = {".txt", ".csv", ".json", ".xml"}


def _local_sample_path(train_path: str,saved_path: str, saved_filename: str) -> str:
    return os.path.join(train_path, saved_path, saved_filename).replace("\\", "/")


def _sample_object_name(saved_path: str, saved_filename: str) -> str:
    return f"{str(saved_path or '').strip('/').replace('\\', '/')}/{saved_filename}".strip("/")


def _ensure_local_dir(saved_path: str) -> str:
    local_dir = os.path.join(settings.FILESPACE_ROOT_PATH, saved_path).replace("\\", "/")
    os.makedirs(local_dir, exist_ok=True)
    return local_dir


def _write_local_file(saved_path: str, saved_filename: str, data: bytes) -> str:
    local_dir = _ensure_local_dir(saved_path)
    full_path = os.path.join(local_dir, saved_filename).replace("\\", "/")
    with open(full_path, "wb") as f:
        f.write(data)
    return full_path


def _sidecar_filename(saved_filename: str, suffix: str) -> str:
    stem = os.path.splitext(saved_filename)[0]
    normalized_suffix = suffix if str(suffix).startswith(".") else f".{suffix}"
    return f"{stem}{normalized_suffix}"


def sample_sidecar_names(saved_filename: str):
    return [
        _sidecar_filename(saved_filename, ".xml"),
        _sidecar_filename(saved_filename, ".json"),
    ]


def _normalize_filesize(size_bytes: int) -> str:
    try:
        return str(min(max(int(size_bytes or 0), 0), 9999999999))
    except Exception:
        return "0"


def _image_wh_from_bytes(data: bytes):
    try:
        with Image.open(io.BytesIO(data)) as img:
            width, height = img.size
        return f"{width}x{height}"
    except UnidentifiedImageError:
        return None
    except Exception:
        return None


def _extract_sample_index(saved_filename: str) -> int | None:
    stem = os.path.splitext(os.path.basename(str(saved_filename or "")))[0]
    tail = stem.rsplit("-", 1)[-1]
    if tail.isdigit():
        return int(tail)
    return None


def _next_sample_index(dataset_id: str) -> int:
    max_saved_filename = AiDatasetSample.objects.filter(dataset_id=dataset_id).aggregate(
        max_saved_filename=Max("saved_filename")
    )["max_saved_filename"]
    current_index = _extract_sample_index(max_saved_filename)
    return (current_index or 0) + 1


def _build_saved_filename(dataset, original_filename: str) -> str:
    ext = os.path.splitext(original_filename or "")[1].lower()
    index = _next_sample_index(dataset.id)
    return f"{dataset.dataset_code}-{index:08}{ext}"


def detect_original_type(filename: str) -> str:
    ext = os.path.splitext(filename or "")[1].lower()
    if ext in IMAGE_EXTENSIONS:
        return "01"
    if ext in VIDEO_EXTENSIONS:
        return "02"
    if ext in AUDIO_EXTENSIONS:
        return "03"
    if ext == ".zip":
        return "04"
    if ext in TEXT_EXTENSIONS:
        return "05"
    return "99"


def dataset_type_accepts_extension(dataset_type: str, extension: str) -> bool:
    ext = str(extension or "").lower()
    if dataset_type == "02":
        return ext in IMAGE_EXTENSIONS
    if dataset_type == "03":
        return ext in VIDEO_EXTENSIONS
    if dataset_type == "04":
        return ext in AUDIO_EXTENSIONS
    if dataset_type == "05":
        return ext in TEXT_EXTENSIONS
    return ext in IMAGE_EXTENSIONS | VIDEO_EXTENSIONS | AUDIO_EXTENSIONS | TEXT_EXTENSIONS


def create_dataset_sample_record(
    dataset,
    original_filename: str,
    saved_filename: str,
    size_bytes: int,
    annotation_content=None,
    custom1=None,
    custom2=None,
    custom3=None,
    status: str = "00",
):
    annotation_fields = split_dataset_sample_annotation_content(annotation_content)
    return AiDatasetSample.objects.create(
        id=str(uuid.uuid4()),
        dataset_id=dataset.id,
        saved_path=dataset.dataset_path,
        original_filename=original_filename,
        saved_filename=saved_filename,
        saved_filesize=_normalize_filesize(size_bytes),
        rectangle=annotation_fields["rectangle"],
        polygon=annotation_fields["polygon"],
        status=status,
        create_time=datetime.now(),
        custom1=custom1,
        custom2=custom2,
        custom3=custom3,
    )


def save_bytes_to_dataset(
    dataset,
    original_filename: str,
    data: bytes,
    content_type: str | None = None,
    annotation_content=None,
    custom1=None,
    custom2=None,
    custom3=None,
):
    saved_filename = _build_saved_filename(dataset, original_filename)
    object_name = _sample_object_name(dataset.dataset_path, saved_filename)
    actual_content_type = content_type or mimetypes.guess_type(original_filename)[0] or "application/octet-stream"
    minio_storage.upload_bytes(DATASET_BUCKET, object_name, data, content_type=actual_content_type)
    if custom1 is None and os.path.splitext(saved_filename)[1].lower() in IMAGE_EXTENSIONS:
        custom1 = _image_wh_from_bytes(data)
    sample = create_dataset_sample_record(
        dataset=dataset,
        original_filename=os.path.basename(original_filename),
        saved_filename=saved_filename,
        size_bytes=len(data),
        annotation_content=annotation_content,
        custom1=custom1,
        custom2=custom2,
        custom3=custom3,
    )
    return sample


def save_uploaded_file_to_dataset(dataset, uploaded_file):
    file_bytes = uploaded_file.read()
    uploaded_file.seek(0)
    return save_bytes_to_dataset(
        dataset=dataset,
        original_filename=uploaded_file.name,
        data=file_bytes,
        content_type=uploaded_file.content_type,
    )


def copy_minio_object_to_dataset(
    dataset,
    source_bucket: str,
    source_object: str,
    original_filename: str,
    size_bytes: int | None = None,
    custom2=None,
    custom3=None,
):
    saved_filename = _build_saved_filename(dataset, original_filename)
    target_object = _sample_object_name(dataset.dataset_path, saved_filename)
    minio_storage.copy_object(source_bucket, source_object, DATASET_BUCKET, target_object)
    actual_size = size_bytes
    if actual_size is None:
        try:
            actual_size = int(minio_storage.get_object_info(DATASET_BUCKET, target_object).get("size") or 0)
        except Exception:
            actual_size = 0
    custom1 = None
    if os.path.splitext(saved_filename)[1].lower() in IMAGE_EXTENSIONS:
        try:
            custom1 = _image_wh_from_bytes(minio_storage.download_bytes(DATASET_BUCKET, target_object))
        except Exception:
            custom1 = None
    return create_dataset_sample_record(
        dataset=dataset,
        original_filename=os.path.basename(original_filename),
        saved_filename=saved_filename,
        size_bytes=actual_size or 0,
        custom1=custom1,
        custom2=custom2,
        custom3=custom3,
    )


def dataset_sample_source_exists(dataset_id: str, source_marker: str | None) -> bool:
    if not source_marker:
        return False
    return AiDatasetSample.objects.filter(dataset_id=dataset_id, custom2=source_marker).exists()


def copy_minio_object_to_dataset_result(
    dataset,
    source_bucket: str,
    source_object: str,
    original_filename: str,
    size_bytes: int | None = None,
    source_marker: str | None = None,
    skip_existing: bool = False,
):
    if skip_existing and dataset_sample_source_exists(dataset.id, source_marker):
        return {
            "imported": [],
            "skipped": [{"name": os.path.basename(original_filename), "reason": "duplicate_source"}],
            "failed": [],
        }
    sample = copy_minio_object_to_dataset(
        dataset=dataset,
        source_bucket=source_bucket,
        source_object=source_object,
        original_filename=original_filename,
        size_bytes=size_bytes,
        custom2=source_marker,
    )
    return {"imported": [sample], "skipped": [], "failed": []}


def import_zip_bytes_to_dataset_result(
    dataset,
    zip_bytes: bytes,
    source_marker_prefix: str | None = None,
    skip_existing: bool = False,
):
    imported = []
    skipped = []
    failed = []
    with zipfile.ZipFile(io.BytesIO(zip_bytes), "r") as zip_file:
        for member in zip_file.infolist():
            if member.is_dir():
                continue
            member_name = member.filename.replace("\\", "/")
            filename = os.path.basename(member_name)
            if not filename or filename.startswith("."):
                continue
            ext = os.path.splitext(filename)[1].lower()
            if not dataset_type_accepts_extension(dataset.dataset_type, ext):
                failed.append({"name": filename, "reason": "type_mismatch"})
                continue
            source_marker = f"{source_marker_prefix}:zip:{member_name}" if source_marker_prefix else None
            if skip_existing and dataset_sample_source_exists(dataset.id, source_marker):
                skipped.append({"name": filename, "reason": "duplicate_source"})
                continue
            try:
                with zip_file.open(member, "r") as fp:
                    data = fp.read()
                imported.append(save_bytes_to_dataset(dataset, filename, data, custom2=source_marker))
            except Exception as exc:
                failed.append({"name": filename, "reason": str(exc)})
    return {"imported": imported, "skipped": skipped, "failed": failed}


def import_zip_bytes_to_dataset(dataset, zip_bytes: bytes, source_marker_prefix: str | None = None, skip_existing: bool = False):
    result = import_zip_bytes_to_dataset_result(
        dataset=dataset,
        zip_bytes=zip_bytes,
        source_marker_prefix=source_marker_prefix,
        skip_existing=skip_existing,
    )
    return result["imported"]


def extract_video_frames_to_dataset_result(
    dataset,
    video_bytes: bytes,
    original_filename: str,
    interval_seconds: int = 1,
    max_frames: int = 200,
    image_format: str = "jpg",
    source_marker_prefix: str | None = None,
    skip_existing: bool = False,
):
    if cv2 is None:
        raise RuntimeError("opencv-python 未安装，无法处理视频")

    video_ext = os.path.splitext(original_filename or "")[1].lower() or ".mp4"
    temp_path = None
    cap = None
    imported = []
    skipped = []
    failed = []
    try:
        with NamedTemporaryFile(delete=False, suffix=video_ext) as temp_file:
            temp_file.write(video_bytes)
            temp_path = temp_file.name

        cap = cv2.VideoCapture(temp_path)
        if not cap.isOpened():
            raise RuntimeError("视频文件无法打开")

        fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
        step_frames = max(int(round(interval_seconds * fps)), 1)
        current_frame = 0
        video_stem = os.path.splitext(os.path.basename(original_filename))[0]
        suffix = "jpg" if image_format == "jpeg" else image_format
        encode_ext = ".jpg" if image_format in {"jpg", "jpeg"} else ".png"

        while len(imported) < max_frames:
            ok, frame = cap.read()
            if not ok:
                break
            if current_frame % step_frames == 0:
                source_marker = (
                    f"{source_marker_prefix}:frame:{current_frame:08}"
                    if source_marker_prefix
                    else None
                )
                if skip_existing and dataset_sample_source_exists(dataset.id, source_marker):
                    skipped.append(
                        {
                            "name": f"{video_stem}_{current_frame:08}.{suffix}",
                            "reason": "duplicate_source",
                        }
                    )
                    current_frame += 1
                    continue
                ok_encode, encoded = cv2.imencode(encode_ext, frame)
                if ok_encode:
                    frame_name = f"{video_stem}_{current_frame:08}.{suffix}"
                    imported.append(
                        save_bytes_to_dataset(
                            dataset,
                            frame_name,
                            encoded.tobytes(),
                            custom2=source_marker,
                        )
                    )
                else:
                    failed.append(
                        {
                            "name": f"{video_stem}_{current_frame:08}.{suffix}",
                            "reason": "frame_encode_failed",
                        }
                    )
            current_frame += 1
    finally:
        if cap is not None:
            cap.release()
        if temp_path and os.path.exists(temp_path):
            os.unlink(temp_path)
    return {"imported": imported, "skipped": skipped, "failed": failed}


def extract_video_frames_to_dataset(
    dataset,
    video_bytes: bytes,
    original_filename: str,
    interval_seconds: int = 1,
    max_frames: int = 200,
    image_format: str = "jpg",
    source_marker_prefix: str | None = None,
    skip_existing: bool = False,
):
    result = extract_video_frames_to_dataset_result(
        dataset=dataset,
        video_bytes=video_bytes,
        original_filename=original_filename,
        interval_seconds=interval_seconds,
        max_frames=max_frames,
        image_format=image_format,
        source_marker_prefix=source_marker_prefix,
        skip_existing=skip_existing,
    )
    return result["imported"]


def ensure_sample_local_copy(train_path: str, saved_path: str, saved_filename: str):
    local_path = _local_sample_path(train_path, saved_path, saved_filename)
    if os.path.exists(local_path):
        return local_path
    object_name = _sample_object_name(saved_path, saved_filename)
    _ensure_local_dir(saved_path)
    minio_storage.download_file(DATASET_BUCKET, object_name, local_path)
    return local_path


def ensure_sidecar_local_copy(saved_path: str, saved_filename: str, suffix: str):
    sidecar_name = _sidecar_filename(saved_filename, suffix)
    local_path = _local_sample_path(saved_path, sidecar_name)
    if os.path.exists(local_path):
        return local_path
    object_name = _sample_object_name(saved_path, sidecar_name)
    if not minio_storage.object_exists(DATASET_BUCKET, object_name):
        return local_path
    _ensure_local_dir(saved_path)
    minio_storage.download_file(DATASET_BUCKET, object_name, local_path)
    return local_path


def download_sidecar_bytes(saved_path: str, saved_filename: str, suffix: str):
    sidecar_name = _sidecar_filename(saved_filename, suffix)
    object_name = _sample_object_name(saved_path, sidecar_name)
    try:
        if minio_storage.object_exists(DATASET_BUCKET, object_name):
            return minio_storage.download_bytes(DATASET_BUCKET, object_name)
    except Exception:
        pass

    local_path = _local_sample_path(saved_path, sidecar_name)
    if os.path.exists(local_path) and os.path.isfile(local_path):
        with open(local_path, "rb") as f:
            return f.read()
    return None


def download_sidecar_text(saved_path: str, saved_filename: str, suffix: str, encoding: str = "utf-8"):
    data = download_sidecar_bytes(saved_path, saved_filename, suffix)
    if data is None:
        return None
    return data.decode(encoding, errors="replace")


def upload_sidecar_bytes(saved_path: str, saved_filename: str, suffix: str, data: bytes, content_type: str | None = None):
    sidecar_name = _sidecar_filename(saved_filename, suffix)
    object_name = _sample_object_name(saved_path, sidecar_name)
    actual_content_type = content_type or mimetypes.guess_type(sidecar_name)[0] or "application/octet-stream"
    minio_storage.upload_bytes(DATASET_BUCKET, object_name, data, content_type=actual_content_type)
    return sidecar_name


def upload_sidecar_text(saved_path: str, saved_filename: str, suffix: str, text: str, content_type: str | None = None):
    actual_content_type = content_type
    if actual_content_type is None and str(suffix).lower().endswith("json"):
        actual_content_type = "application/json"
    return upload_sidecar_bytes(
        saved_path=saved_path,
        saved_filename=saved_filename,
        suffix=suffix,
        data=(text or "").encode("utf-8"),
        content_type=actual_content_type,
    )


def delete_sample_storage(saved_path: str, saved_filename: str):
    candidates = [saved_filename] + sample_sidecar_names(saved_filename)
    failed_objects = []
    for filename in candidates:
        object_name = _sample_object_name(saved_path, filename)
        try:
            if minio_storage.object_exists(DATASET_BUCKET, object_name):
                minio_storage.delete_object(DATASET_BUCKET, object_name)
        except Exception:
            failed_objects.append(object_name)

        local_path = os.path.join(settings.FILESPACE_ROOT_PATH, saved_path, filename).replace("\\", "/")
        if os.path.exists(local_path) and os.path.isfile(local_path):
            try:
                os.remove(local_path)
            except Exception:
                pass
    return failed_objects