import time
from multi_model_detection.models import OperationRecord
from app01.models import UserInfo
from django.utils import timezone
from django.contrib.auth.decorators import login_required
from django.shortcuts import render
from PAD_predict.views import get_pest_model
@login_required
def multi_model_detect(request):
    return render(request, 'multi_model_detection.html', {'active_page': 'multi_model_detection'})

import json
import cv2
import numpy as np
import os
import base64
import tempfile
from django.contrib.auth.decorators import login_required
from django.views.decorators.http import require_http_methods
from django.http import JsonResponse, HttpResponse, StreamingHttpResponse
from django.views.decorators.csrf import csrf_exempt
from django.core.files.storage import FileSystemStorage
from ultralytics import YOLO

from djangoProject import settings
from PAD_predict.models import PestDetectionRecord  # 使用新的模型

MODEL_PATHS = {
    "YOLO11-1": r'D:\eggplantshapeBS\djangoProject\PAD_train\src\src\results\exp4\weights\best.pt',
    "YOLO11-2": r'D:\eggplantshapeBS\djangoProject\PAD_train\src\src\results\exp4\weights\last.pt',
    "model3": r"D:\eggplantshapeBS\djangoProject\PAD_train\src\src\results\exp4\weights\last.pt",
    "YOLOv8": r'D:\eggplantshapeBS\djangoProject\PAD_train\src\src\results\exp2\weights\best.pt',
    "model4": r'D:\eggplantshapeBS\djangoProject\PAD_train\src\yolo11n.pt',
}

# 加载所有模型
models = {name: YOLO(path) for name, path in MODEL_PATHS.items()}

@login_required
@require_http_methods(["GET"])
def get_models(request):
    # 获取所有模型名称
    model_names = list(MODEL_PATHS.keys())
    return JsonResponse({'models': model_names})


@login_required
@require_http_methods(["POST"])
@csrf_exempt
def upload_and_detect_pest(request):
    image_file = request.FILES.get('image')
    selected_models = request.POST.getlist('model')

    if not image_file:
        return JsonResponse({'error': 'No image file provided'}, status=400)

    fs = FileSystemStorage()

    save_name = f"uploads/{timezone.now().strftime('%Y%m%d%H%M%S')}_{image_file.name}"
    filename = fs.save(save_name, image_file)
    image_path = fs.path(filename)

    image = cv2.imread(image_path)
    results = []
    start_time = time.time()

    for model_name in selected_models:
        result_data = {'model': model_name}
        if model_name in models:
            try:
                result = models[model_name].predict(image)
                if result:
                    annotated_image = result[0].plot()
                    _, buffer = cv2.imencode('.jpg', annotated_image)
                    result_data['resultImage'] = base64.b64encode(buffer).decode('utf-8')
                    result_data['count'] = len(result[0].boxes)
                else:
                    result_data['error'] = "No detections found"
            except Exception as e:
                result_data['error'] = str(e)
        else:
            result_data['error'] = "Model not found"
        results.append(result_data)

    end_time = time.time()
    duration = round(end_time - start_time, 3)

    try:
        user_info_instance = UserInfo.objects.get(user_id=request.user.id)

        OperationRecord.objects.create(
            user=user_info_instance,
            operation_type="IMAGE",
            description=f"使用模型 {', '.join(selected_models)} 进行检测。耗时: {duration}s",
            result_path=filename
        )
    except UserInfo.DoesNotExist:

        logger.error(f"未找到用户 ID 为 {request.user.id} 的 UserInfo 记录")
    except Exception as e:
        logger.error(f"保存操作记录失败: {str(e)}")


    return JsonResponse({
        'results': results,
        'inference_time': duration,
        'total_count': sum([r.get('count', 0) for r in results])
    })

import logging

logger = logging.getLogger(__name__)

@login_required
@require_http_methods(["POST"])
@csrf_exempt
def upload_video_and_detect_pest(request):
    video_file = request.FILES.get('video')
    selected_models = request.POST.getlist('model')

    if not video_file or not selected_models:
        logger.error("No video file or models provided")
        return JsonResponse({'error': 'No video file or models provided'}, status=400)

    with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(video_file.name)[1]) as temp_file:
        temp_video_path = temp_file.name
        for chunk in video_file.chunks():
            temp_file.write(chunk)

    logger.info(f"Video saved to temporary file: {temp_video_path}")

    cap = cv2.VideoCapture(temp_video_path)
    if not cap.isOpened():
        logger.error("Failed to open video file")
        return JsonResponse({'error': 'Failed to open video file'}, status=400)

    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    logger.info(f"Video properties: FPS={fps}, Width={width}, Height={height}")

    response_data = {}

    for model_name in selected_models:
        if model_name not in models:
            response_data[model_name] = {'error': 'Model not found'}
            continue

        output_video_path = tempfile.mktemp(suffix='.mp4')
        fourcc = cv2.VideoWriter_fourcc(*'avc1')
        out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))

        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
        while True:
            ret, frame = cap.read()
            if not ret:
                break

            results = models[model_name].predict(frame)
            if results:
                result = results[0]
                annotated_frame = result.plot()
                out.write(annotated_frame)

        out.release()

        with open(output_video_path, 'rb') as video_file:
            video_bytes = video_file.read()
        video_base64 = base64.b64encode(video_bytes).decode('utf-8')
        response_data[model_name] = video_base64

        os.remove(output_video_path)

    cap.release()
    os.remove(temp_video_path)

    logger.info("Video processing completed")
    response = HttpResponse(json.dumps(response_data), content_type='application/json')
    response['Content-Disposition'] = 'attachment; filename="detected_videos.json"'
    return response

@login_required
@csrf_exempt
@require_http_methods(["GET"])
def camera_detect_pest(request):
    selected_models = request.GET.get('model')

    if not selected_models:
        return JsonResponse({'error': 'No model names provided'}, status=400)

    selected_models = selected_models.split(',')

    def generate_frames():
        cap = cv2.VideoCapture(0)
        if not cap.isOpened():
            yield f"data: {{'error': 'Failed to open camera'}}\n\n"
            return

        try:
            while True:
                ret, frame = cap.read()
                if not ret:
                    break

                model_results = {}
                for model_name in selected_models:
                    try:
                        model = models[model_name]
                        results = model.predict(frame)
                        if results:
                            annotated_frame = results[0].plot()
                            _, buffer = cv2.imencode('.jpg', annotated_frame)
                            model_results[model_name] = f"data:image/jpeg;base64,{base64.b64encode(buffer).decode('utf-8')}"
                        else:
                            model_results[model_name] = "No detections found"
                    except Exception as e:
                        model_results[model_name] = f"Error: {str(e)}"

                yield f"data: {json.dumps(model_results)}\n\n"
        finally:
            cap.release()

    return StreamingHttpResponse(generate_frames(), content_type='text/event-stream')