#!/bin/sh

# RENAME THIS FILE AS INIT.SH, THEN RUN IT AS `bash init.sh`

read -p "Do you want to set up your Python environment in $(pwd)?
[y/n]: " confirm
if [ "$confirm" != "y" ]; then
  echo "Setup aborted."
  exit 1
fi

mkdir python-ai

cd python-ai

python3 -m venv venv
source venv/bin/activate

pip install ollama pydantic

echo "
from ollama import chat
from pydantic import BaseModel

# From https://ollama.com/blog/structured-outputs

class Pet(BaseModel):
  name: str
  animal: str
  age: int
  color: str | None
  favorite_toy: str | None

class PetList(BaseModel):
  pets: list[Pet]

response = chat(
  messages=[
    {
      'role': 'user',
      'content': ''' //You can change this prompt to test different inputs
        I have two pets.
        A cat named Luna who is 5 years old and loves playing with yarn. She has grey fur.
        I also have a 2 year old black cat named Loki who loves tennis balls.
      ''',
    }
  ],
  model='gemma3:1b',
  format=PetList.model_json_schema(),
)

pets = PetList.model_validate_json(response.message.content)
print(pets)
" > example.py

ollama ls | grep "gemma3:1b"

if [ $? -eq 0 ]; then
  echo "The gemma3:1b model is installed."
else
    read -p "Download the gemma3:1b model now? You'll need this to run the example.
    [y/n]: " confirm
    if [ "$confirm" != "y" ]; then
      echo "Gemma download aborted."
      exit 1
    else
        ollama pull gemma3:1b
    fi
fi



read -p "Environment created successfully
Run example.py script?
[y/n]: " confirm
if [ "$confirm" != "y" ]; then
  echo "Example script aborted."
  exit 1
else
    echo "Running simulated query..."
    sleep 1
    TYPEDTEXT="USER: I have two pets. A cat named Luna who is 5 years old and loves playing with yarn. She has grey fur. I also have a 2 year old black cat named Loki who loves tennis balls. Return structured data about them"
    for ((i = 0; i < ${#TYPEDTEXT}; i++)); do
        echo -n "${TYPEDTEXT:$i:1}"
        sleep 0.02
    done
    echo
    python3 example.py
fi