from celery import shared_task
from .models import Track
from django.core.files import File
from django.core.files.uploadedfile import InMemoryUploadedFile
from django.core.files.base import ContentFile

import os
import dotenv

import json
from requests import get, post
from requests.exceptions import ConnectionError
from PIL import Image
from io import BytesIO

dotenv.load_dotenv()

@shared_task
def get_banter(id):
    """
        Requests an AI model for some banter.
        Currently expects an ollama instance to be running,
        but might be possible to use OpenAI API in future.
    """
    track = Track.objects.get(pk=id)

    try:
        get(os.getenv('AI_ENDPOINT'))
    except ConnectionError:
        track.banter = "Error"
        track.banter_done = True
        track.save()
        return f"{track}: AI endpoint not available."

    try:
        r = post(
            os.getenv('AI_ENDPOINT'),
            json={
                'model': os.getenv('AI_MODEL'),
                'prompt': str(track)
            },
            stream=True
        )
        r.raise_for_status()

        for line in r.iter_lines():
            body = json.loads(line)
            if 'error' in body:
                raise Exception(body['error'])
            track.banter += body.get('response', '')

            if body.get('done', False):
                track.banter_done = True
                track.save()
                return f"{track}: {body['context']}"
    except Exception as e:
        track.banter_done = True
        track.save()
        return f"{track}: {e}"

@shared_task
def get_and_dither_image(url, instance, attribute):
    """Expects a url, and a Django model cls with an ImageField attribute"""
    print(f"Parsing {url} for {instance} with attribute {attribute}")
    r = get(url)
    print(r.content)
    i = Image.open(BytesIO(r.content))
    i = i.resize(size=(256, 256)).convert("P").quantize(colors=32)
    buffer = BytesIO()
    i.save(fp=buffer, format='PNG')
    file = InMemoryUploadedFile(ContentFile(buffer.getvalue()), None, instance.name + ".png", "image/png", buffer.tell, "utf-8")
    instance.image = file
    instance.save()