"""finetune.py, M15: submit a fine-tuning job and use the result.

Fine-tuning needs a provider that hosts training (GPUs cost money/time). The simplest hosted
path is OpenAI's fine-tuning API: upload your JSONL, start a job, wait, then call your new
model by its id. (Anthropic doesn't offer a simple public fine-tuning API, so this example
uses OpenAI, a different provider and key. The free/local alternative is at the bottom.)

Setup:  pip install openai   and put OPENAI_API_KEY in .env
Run:    python finetune.py train.jsonl          # start a job
        python finetune.py status ftjob-xxxxx   # check / get your model id
"""

import os
import sys
from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()
client = OpenAI()                                  # reads OPENAI_API_KEY
BASE_MODEL = "gpt-4o-mini-2024-07-18"              # a fine-tunable base model


def submit(training_file_path):
    """Upload the dataset and start a fine-tuning job. Returns the job id."""
    uploaded = client.files.create(
        file=open(training_file_path, "rb"), purpose="fine-tune",
    )
    job = client.fine_tuning.jobs.create(training_file=uploaded.id, model=BASE_MODEL)
    return job.id


def check(job_id):
    """Return (status, fine_tuned_model). status goes queued -> running -> succeeded."""
    job = client.fine_tuning.jobs.retrieve(job_id)
    return job.status, job.fine_tuned_model


def use(fine_tuned_model, user_message):
    """Call YOUR fine-tuned model, note it needs no long system prompt; the style is baked in."""
    resp = client.chat.completions.create(
        model=fine_tuned_model,
        messages=[{"role": "user", "content": user_message}],
    )
    return resp.choices[0].message.content


if __name__ == "__main__":
    if len(sys.argv) >= 3 and sys.argv[1] == "status":
        status, model = check(sys.argv[2])
        print(f"status: {status}\nfine_tuned_model: {model}")
        if model:
            print("\nTry it:  use(model, 'my latte was cold')  → an on-brand reply, no prompt needed.")
    else:
        path = sys.argv[1] if len(sys.argv) > 1 else "train.jsonl"
        job_id = submit(path)
        print(f"Started fine-tuning job: {job_id}")
        print(f"Check progress with:  python finetune.py status {job_id}")
        print("(Training takes minutes to hours and costs a few dollars on a small dataset.)")

# --- Free / local alternative (open models, no per-token bill) ----------------
# To fine-tune an OPEN model on your own machine/GPU instead of a hosted API, use
# Hugging Face: `transformers` + `peft` (LoRA), or the friendly `unsloth` wrapper.
# LoRA trains a few small "adapter" layers instead of the whole model, so it's cheap
# and fits on a modest GPU. Same dataset idea (examples of input -> ideal output);
# you load a base model (e.g. Llama/Gemma), train the adapter, then run it with Ollama
# (M13) or transformers. (Needs a GPU to be practical, outline only here.)