from fastapi import UploadFile, Form
from fastapi.responses import JSONResponse
from pipeline_setup import pipe
from utils.image_processing import encode_image_base64
from utils.video_processing import split_video_into_segments, extract_motion_key_frames, extract_audio_from_video
from utils.audio_transcription import transcribe_audio
import asyncio
import mimetypes
from concurrent.futures import ThreadPoolExecutor

# async def video_query(file: UploadFile, question: str = Form(...)):
#     """
#     API endpoint to process a video file with the user's query.
#     """
#     try:
#         print("Processing video...")

#         # Validate file type
#         if file.content_type not in ["video/mp4", "video/avi", "video/mkv"]:
#             return JSONResponse({"query": question, "error": "Unsupported video file type."})

#         # Start overall timer
#         overall_start_time = time.time()

#         # Save the uploaded video to a temporary file
#         print("Reading video...")
#         video_data = await file.read()
#         temp_video_path = "/tmp/temp_video.mp4"
#         with open(temp_video_path, "wb") as temp_video_file:
#             temp_video_file.write(video_data)
#         print(f"Temp video saved to: {temp_video_path}")

#         # Record the time after reading the video
#         video_reading_time = time.time()

#         # Split the video into segments
#         print("Splitting video...")
#         segments = split_video_into_segments(temp_video_path, segment_duration=30)
#         print(f"Video split into {len(segments)} segments.")

#         aggregated_responses = []
#         segment_timings = []

#         for i, segment_path in enumerate(segments):
#             print(f"Processing segment {i+1}/{len(segments)}: {segment_path}")

#             # Start timing for the segment
#             segment_start_time = time.time()

#             # Extract key frames
#             frame_start_time = time.time()
#             imgs = extract_motion_key_frames(segment_path, max_frames=50, sigma_multiplier=2)
#             frame_time = time.time()

#             # Extract audio and transcribe
#             audio_start_time = time.time()
#             audio_path = extract_audio_from_video(segment_path)
#             transcribed_text = transcribe_audio(audio_path)
#             audio_time = time.time()

#             # Combine transcribed text with the query
#             combined_query = f"Audio Transcript: {transcribed_text}\n{question}"

#             # Prepare content for the pipeline
#             question_with_frames = ""
#             for j, img in enumerate(imgs):
#                 question_with_frames += f"Frame{j+1}: {{IMAGE_TOKEN}}\n"
#             question_with_frames += combined_query

#             content = [{"type": "text", "text": question_with_frames}]
#             for img in imgs:
#                 content.append({
#                     "type": "image_url",
#                     "image_url": {
#                         "max_dynamic_patch": 1,
#                         "url": f"data:image/jpeg;base64,{encode_image_base64(img)}"
#                     }
#                 })

#             # Query the model
#             inference_start_time = time.time()
#             messages = [dict(role="user", content=content)]
#             response = await asyncio.to_thread(pipe, messages)
#             inference_time = time.time()

#             # Aggregate response
#             aggregated_responses.append(response.text)

#             # Calculate timing for the segment
#             segment_timings.append({
#                 "segment_index": i + 1,
#                 "segment_processing_time": inference_time - segment_start_time,
#                 "frame_extraction_time": frame_time - frame_start_time,
#                 "audio_extraction_time": audio_time - audio_start_time,
#                 "model_inference_time": inference_time - inference_start_time
#             })

#             print(f"transcription: {transcribed_text}")
#             # print(f"content: {content}")

#         overall_end_time = time.time()

#         # Aggregate total timings
#         total_timings = {
#             "video_reading_time": video_reading_time - overall_start_time,
#             "total_segments": len(segments),
#             "total_processing_time": overall_end_time - overall_start_time,
#             "segment_details": segment_timings
#         }

#         return JSONResponse({
#             "question": question,
#             "responses": aggregated_responses,
#             "timings": total_timings,
#         })
#     except Exception as e:
#         return JSONResponse({"query": question, "error": str(e)})


async def video_query(video_path: str, question: str):
    """
    API endpoint to process a video file with the user's query.
    """
    try:
        print("Processing video...")

        if not video_path or not isinstance(video_path, str):
            return {"query": question, "error": "No video file provided or invalid file input."}

        # Determine the file type using the file extension
        file_type, _ = mimetypes.guess_type(video_path)
        if file_type is None or not file_type.startswith("video/"):
            return {"query": question, "error": "Unsupported video file type."}

        # Log the video path
        print(f"Video path: {video_path}")

        # Split the video into segments
        print("Splitting video...")
        segments = split_video_into_segments(video_path, segment_duration=30)
        print(f"Video split into {len(segments)} segments.")

        aggregated_responses = []
        segment_timings = []

        for i, segment_path in enumerate(segments):
            print(f"Processing segment {i+1}/{len(segments)}: {segment_path}")

            # Extract key frames
            imgs = extract_motion_key_frames(segment_path, max_frames=50, sigma_multiplier=2)

            # Extract audio and transcribe
            audio_path = extract_audio_from_video(segment_path)
            transcribed_text = transcribe_audio(audio_path)

            # Combine transcribed text with the query
            combined_query = f"Audio Transcript: {transcribed_text}\n{question}"

            # Prepare content for the pipeline
            question_with_frames = ""
            for j, img in enumerate(imgs):
                question_with_frames += f"Frame{j+1}: {{IMAGE_TOKEN}}\n"
            question_with_frames += combined_query

            content = [{"type": "text", "text": question_with_frames}]
            for img in imgs:
                content.append({
                    "type": "image_url",
                    "image_url": {
                        "max_dynamic_patch": 1,
                        "url": f"data:image/jpeg;base64,{encode_image_base64(img)}"
                    }
                })

            # Query the model
            messages = [dict(role="user", content=content)]
            response = await asyncio.to_thread(pipe, messages)

            # Aggregate response
            aggregated_responses.append(response.text)

        return {
            "question": question,
            "responses": aggregated_responses,
        }
    except Exception as e:
        return {"query": question, "error": str(e)}