from celery import shared_task
from .models import Track

import dotenv

import os
import requests

import time
import json
from random import choice

dotenv.load_dotenv()

@shared_task
def get_banter(id):
    track = Track.objects.get(pk=id)
    r = requests.post(
        os.getenv('AI_ENDPOINT'),
        #'http://localhost:11434/api/generate',
        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)
        response_part = body.get('response', '')
        # the response streams one token at a time, print that as we receive it
        print(response_part, end='', flush=True)
        track.banter += response_part
        track.save()

        if 'error' in body:
            raise Exception(body['error'])

        if body.get('done', False):
            track.banter_done = True
            track.save()
            return body['context']
@shared_task
def get_banter_test(id):
    letters = ["a","b","c","d","e","q","w", " "]
    track = Track.objects.get(pk=id)
    time.sleep(20)
    track.banter = "Lorem "
    track.save()
    time.sleep(.5)

    for i in range(100):
        track.banter += choice(letters)
        track.save()
        time.sleep(.3)
    track.banter_done = True
    track.save()