"""
fine_tune_model.py — Fine-Tuning Workflow for Azure OpenAI
==========================================================
End-to-end fine-tuning pipeline:
  1. Prepare training data in JSONL format
  2. Upload the training file
  3. Create a fine-tuning job
  4. Monitor job progress
  5. Deploy the fine-tuned model

Prerequisites:
    pip install openai azure-identity python-dotenv rich

Usage:
    python fine_tune_model.py
"""

import json
import os
import sys
import time
from dotenv import load_dotenv
from openai import AzureOpenAI
from azure.identity import DefaultAzureCredential, get_bearer_token_provider
from rich.console import Console
from rich.progress import Progress

load_dotenv()
console = Console()


def get_env(key: str) -> str:
    value = os.getenv(key)
    if not value:
        console.print(f"[red]ERROR:[/red] '{key}' not set. See .env.example.")
        sys.exit(1)
    return value


def get_client() -> AzureOpenAI:
    endpoint = get_env("AZURE_OPENAI_ENDPOINT")
    api_version = get_env("AZURE_OPENAI_API_VERSION")
    api_key = os.getenv("AZURE_OPENAI_API_KEY")

    if api_key:
        return AzureOpenAI(azure_endpoint=endpoint, api_key=api_key, api_version=api_version)

    token_provider = get_bearer_token_provider(
        DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
    )
    return AzureOpenAI(azure_endpoint=endpoint, azure_ad_token_provider=token_provider, api_version=api_version)


# ---------------------------------------------------------------------------
# Step 1 — Prepare training data
# ---------------------------------------------------------------------------
TRAINING_EXAMPLES = [
    {
        "messages": [
            {"role": "system", "content": "You are an Azure pricing assistant. Answer concisely."},
            {"role": "user", "content": "How much does Azure AI Search cost?"},
            {"role": "assistant", "content": "Azure AI Search Basic tier starts at ~$75.61/month. Free tier available with limits."},
        ]
    },
    {
        "messages": [
            {"role": "system", "content": "You are an Azure pricing assistant. Answer concisely."},
            {"role": "user", "content": "What is the cost of GPT-4o on Azure?"},
            {"role": "assistant", "content": "GPT-4o on Azure OpenAI: ~$2.50/1M input tokens, ~$10.00/1M output tokens (Global Standard)."},
        ]
    },
    {
        "messages": [
            {"role": "system", "content": "You are an Azure pricing assistant. Answer concisely."},
            {"role": "user", "content": "Is Azure AI Foundry free?"},
            {"role": "assistant", "content": "Azure AI Foundry itself has no extra charge. You pay for underlying resources (compute, storage, models)."},
        ]
    },
]


def prepare_training_file(path: str) -> str:
    """Write training examples to a JSONL file."""
    console.print("\n[bold]Step 1:[/bold] Preparing training data...", style="cyan")

    with open(path, "w", encoding="utf-8") as f:
        for example in TRAINING_EXAMPLES:
            f.write(json.dumps(example) + "\n")

    console.print(f"   ✅ Wrote {len(TRAINING_EXAMPLES)} examples to {path}")
    console.print("   ℹ️  For production, you need at least 10 examples (50+ recommended).")
    return path


# ---------------------------------------------------------------------------
# Step 2 — Upload training file
# ---------------------------------------------------------------------------
def upload_file(client: AzureOpenAI, path: str) -> str:
    """Upload the JSONL file to Azure OpenAI."""
    console.print("\n[bold]Step 2:[/bold] Uploading training file...", style="cyan")

    with open(path, "rb") as f:
        result = client.files.create(file=f, purpose="fine-tune")

    console.print(f"   ✅ File uploaded: {result.id} (status: {result.status})")
    return result.id


# ---------------------------------------------------------------------------
# Step 3 — Create fine-tuning job
# ---------------------------------------------------------------------------
def create_fine_tuning_job(client: AzureOpenAI, file_id: str, model: str = "gpt-4o-mini-2024-07-18") -> str:
    """Create a fine-tuning job."""
    console.print("\n[bold]Step 3:[/bold] Creating fine-tuning job...", style="cyan")

    job = client.fine_tuning.jobs.create(
        training_file=file_id,
        model=model,
        hyperparameters={
            "n_epochs": 3,
        },
    )

    console.print(f"   ✅ Job created: {job.id}")
    console.print(f"   Base model: {model}")
    return job.id


# ---------------------------------------------------------------------------
# Step 4 — Monitor job
# ---------------------------------------------------------------------------
def monitor_job(client: AzureOpenAI, job_id: str) -> str:
    """Poll the fine-tuning job until it completes."""
    console.print("\n[bold]Step 4:[/bold] Monitoring fine-tuning job...", style="cyan")

    with Progress(console=console) as progress:
        task = progress.add_task("Fine-tuning...", total=None)

        while True:
            job = client.fine_tuning.jobs.retrieve(job_id)
            status = job.status

            progress.update(task, description=f"Status: {status}")

            if status == "succeeded":
                console.print(f"\n   ✅ Fine-tuning complete!")
                console.print(f"   Fine-tuned model: {job.fine_tuned_model}")
                return job.fine_tuned_model

            if status in ("failed", "cancelled"):
                console.print(f"\n   ❌ Job {status}: {job.error}")
                sys.exit(1)

            time.sleep(30)


# ---------------------------------------------------------------------------
# Step 5 — Test the fine-tuned model
# ---------------------------------------------------------------------------
def test_model(client: AzureOpenAI, model_name: str) -> None:
    """Send a test prompt to the fine-tuned model."""
    console.print("\n[bold]Step 5:[/bold] Testing fine-tuned model...", style="cyan")

    response = client.chat.completions.create(
        model=model_name,
        messages=[
            {"role": "system", "content": "You are an Azure pricing assistant. Answer concisely."},
            {"role": "user", "content": "How much does Azure Blob Storage cost?"},
        ],
        max_tokens=128,
    )

    console.print(f"   🤖 {response.choices[0].message.content}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> None:
    client = get_client()

    console.print("[bold]Azure OpenAI Fine-Tuning Pipeline[/bold]", style="magenta")
    console.print("=" * 60)

    training_file = "training_data.jsonl"
    prepare_training_file(training_file)
    file_id = upload_file(client, training_file)
    job_id = create_fine_tuning_job(client, file_id)
    fine_tuned_model = monitor_job(client, job_id)
    test_model(client, fine_tuned_model)

    # Clean up local file
    os.remove(training_file)

    console.print("\n[bold green]✅ Fine-tuning pipeline complete![/bold green]")


if __name__ == "__main__":
    main()