from celery import shared_task
from .models import Track

import os
import dotenv

import json
from requests import get, post
from requests.exceptions import ConnectionError

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}"