from rest_framework.views import APIView
import tempfile, os
from rest_framework.response import Response

from _utils.resumo_completo_cursor import get_llm_summary_answer_by_cursor
from _utils.utils import DEFAULT_SYSTEM_PROMPT
from .serializer import ResumoPDFSerializer, ResumoCursorSerializer
from _utils.main import get_llm_answer_summary, get_llm_answer_summary_with_embedding
from setup.environment import default_model
from rest_framework.parsers import MultiPartParser
from drf_spectacular.utils import extend_schema


class ResumoView(APIView):
  parser_classes = [MultiPartParser]

  @extend_schema(
    request=ResumoPDFSerializer,
  )
  def post(self, request):
    serializer = ResumoPDFSerializer(data=request.data)
    if serializer.is_valid(raise_exception=True):
        listaPDFs = []
        data = serializer.validated_data
        model = serializer.validated_data.get("model", default_model)
        print('serializer.validated_data: ', serializer.validated_data)

        for file in serializer.validated_data['files']:
            print("file: ", file)
            file.seek(0)
            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file: # Create a temporary file to save the uploaded PDF
                for chunk in file.chunks(): # Write the uploaded file content to the temporary file
                    temp_file.write(chunk)
                temp_file_path = temp_file.name  # Get the path of the temporary file
                listaPDFs.append(temp_file_path)
            # print('listaPDFs: ', listaPDFs)

        system_prompt = data.get("system_prompt", DEFAULT_SYSTEM_PROMPT)
        resposta_llm = get_llm_answer_summary(system_prompt, data["user_message"], listaPDFs, model=model, isIterativeRefinement=data["iterative_refinement"])

        for file in listaPDFs:
            os.remove(file)

        return Response({"resposta": resposta_llm})

class ResumoEmbeddingView(APIView):
  parser_classes = [MultiPartParser]

  @extend_schema(
    request=ResumoPDFSerializer,
  )
  def post(self, request):
    serializer = ResumoPDFSerializer(data=request.data)
    if serializer.is_valid(raise_exception=True):
        listaPDFs = []
        data = serializer.validated_data
        model = serializer.validated_data.get("model", default_model)
        print('serializer.validated_data: ', serializer.validated_data)

        for file in serializer.validated_data['files']:
            file.seek(0)
            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file: # Create a temporary file to save the uploaded PDF
                for chunk in file.chunks(): # Write the uploaded file content to the temporary file
                    temp_file.write(chunk)
                temp_file_path = temp_file.name  # Get the path of the temporary file
                listaPDFs.append(temp_file_path)
        print('listaPDFs: ', listaPDFs)

        system_prompt = data.get("system_prompt", DEFAULT_SYSTEM_PROMPT)
        resposta_llm = get_llm_answer_summary_with_embedding(system_prompt, data["user_message"], listaPDFs, model=model, isIterativeRefinement=data["iterative_refinement"])

        for file in listaPDFs:
            os.remove(file)

        return Response({"resposta": resposta_llm})

class ResumoCompletoCursorView(APIView):
  parser_classes = [MultiPartParser]

  @extend_schema(
    request=ResumoCursorSerializer,
  )
  def post(self, request):
    serializer = ResumoCursorSerializer(data=request.data)
    if serializer.is_valid(raise_exception=True):
        listaPDFs = []
        data = serializer.validated_data
        print('\nserializer.validated_data: ', serializer.validated_data)

        for file in serializer.validated_data['files']:
            file.seek(0)
            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file: # Create a temporary file to save the uploaded PDF
                for chunk in file.chunks(): # Write the uploaded file content to the temporary file
                    temp_file.write(chunk)
                temp_file_path = temp_file.name  # Get the path of the temporary file
                listaPDFs.append(temp_file_path)
        print('listaPDFs: ', listaPDFs)

        resposta_llm = get_llm_summary_answer_by_cursor(data, listaPDFs)

        for file in listaPDFs:
            os.remove(file)

        return Response({"resposta": resposta_llm})