from typing import Dict, List, Union
from openai import OpenAI

import typer
import time
from rich.live import Live
from rich.table import Table
from rich.console import Console
from datetime import datetime
from typing import cast
from openai.types.fine_tuning import FineTuningJob

client = OpenAI()
app = typer.Typer()
console = Console()


def generate_table(jobs: List[FineTuningJob]) -> Table:
    # Sorting the jobs by creation time
    jobs = sorted(jobs, key=lambda x: (cast(FineTuningJob, x)).created_at, reverse=True)
    jobs = cast(List[FineTuningJob], jobs)

    table = Table(
        title="OpenAI Fine Tuning Job Monitoring",
        caption="Automatically refreshes every 5 seconds, press Ctrl+C to exit",
    )

    table.add_column("Job ID", style="dim")
    table.add_column("Status")
    table.add_column("Creation Time", justify="right")
    table.add_column("Completion Time", justify="right")
    table.add_column("Model Name")
    table.add_column("File ID")
    table.add_column("Epochs")
    table.add_column("Base Model")

    for job in jobs:
        status_emoji = {
            "running": "⏳",
            "succeeded": "✅",
            "failed": "❌",
            "cancelled": "🚫",
        }.get(job.status, "❓")

        finished_at = (
            str(datetime.fromtimestamp(job.finished_at)) if job.finished_at else "N/A"
        )

        table.add_row(
            job.id,
            f"{status_emoji} [{status_color(job.status)}]{job.status}[/]",
            str(datetime.fromtimestamp(job.created_at)),
            finished_at,
            job.fine_tuned_model,
            job.training_file,
            str(job.hyperparameters.n_epochs),
            job.model,
        )

    return table


def status_color(status: str) -> str:
    return {"running": "yellow", "succeeded": "green", "failed": "red"}.get(
        status, "white"
    )


def get_jobs(limit: int = 5) -> List[FineTuningJob]:
    return client.fine_tuning.jobs.list(limit=limit).data


def get_file_status(file_id: str) -> str:
    response = client.files.retrieve(file_id)
    return response.status


@app.command(
    name="list",
    help="Monitor the status of the most recent fine-tuning jobs.",
)  # type: ignore[misc]
def watch(
    limit: int = typer.Option(5, help="Limit the number of jobs to monitor"),
    poll: int = typer.Option(5, help="Polling interval in seconds"),
    screen: bool = typer.Option(False, help="Enable or disable screen output"),
) -> None:
    """
    Monitor the status of the most recent fine-tuning jobs.
    """
    jobs = get_jobs(limit=limit)
    with Live(generate_table(jobs), refresh_per_second=2, screen=screen) as live_table:
        while True:
            jobs = get_jobs(limit=limit)
            live_table.update(generate_table(jobs))
            time.sleep(poll)


@app.command(
    help="Create a fine-tuning job from an existing ID.",
)  # type: ignore[misc]
def create_from_id(
    id: str = typer.Argument(..., help="ID of the existing fine-tuning job"),
    model: str = typer.Option("gpt-3.5-turbo", help="Model to use for fine-tuning"),
    n_epochs: int = typer.Option(
        None, help="Number of epochs for fine-tuning", show_default=False
    ),
    batch_size: int = typer.Option(
        None, help="Batch size for fine-tuning", show_default=False
    ),
    learning_rate_multiplier: float = typer.Option(
        None, help="Learning rate multiplier for fine-tuning", show_default=False
    ),
    validation_file_id: str = typer.Option(
        None, help="ID of the uploaded validation file"
    ),
) -> None:
    hyperparameters_dict: Dict[str, Union[int, float, str]] = {}
    if n_epochs is not None:
        hyperparameters_dict["n_epochs"] = n_epochs
    if batch_size is not None:
        hyperparameters_dict["batch_size"] = batch_size
    if learning_rate_multiplier is not None:
        hyperparameters_dict["learning_rate_multiplier"] = learning_rate_multiplier

    with console.status(
        f"[bold green]Creating fine-tuning job from ID {id}...", spinner="dots"
    ):
        job = client.fine_tuning.jobs.create(
            training_file=id,
            model=model,
            hyperparameters=hyperparameters_dict if hyperparameters_dict else None,
            validation_file=validation_file_id if validation_file_id else None,
        )
        console.log(f"[bold green]Fine-tuning job created with ID: {job.id}")
    watch(limit=5, poll=2, screen=False)


@app.command(
    help="Create a fine-tuning job from a file.",
)  # type: ignore[misc]
def create_from_file(
    file: str = typer.Argument(..., help="Path to the file for fine-tuning"),
    model: str = typer.Option("gpt-3.5-turbo", help="Model to use for fine-tuning"),
    poll: int = typer.Option(2, help="Polling interval in seconds"),
    n_epochs: int = typer.Option(
        None, help="Number of epochs for fine-tuning", show_default=False
    ),
    batch_size: int = typer.Option(
        None, help="Batch size for fine-tuning", show_default=False
    ),
    learning_rate_multiplier: float = typer.Option(
        None, help="Learning rate multiplier for fine-tuning", show_default=False
    ),
    validation_file: str = typer.Option(None, help="Path to the validation file"),
    model_suffix: str = typer.Option(None, help="Suffix to identify the model"),
) -> None:
    hyperparameters_dict: Dict[str, Union[int, float, str]] = {}
    if n_epochs is not None:
        hyperparameters_dict["n_epochs"] = n_epochs
    if batch_size is not None:
        hyperparameters_dict["batch_size"] = batch_size
    if learning_rate_multiplier is not None:
        hyperparameters_dict["learning_rate_multiplier"] = learning_rate_multiplier

    with open(file, "rb") as file_buffer:
        response = client.files.create(file=file_buffer, purpose="fine-tune")

    file_id = response.id

    validation_file_id = None
    if validation_file:
        with open(validation_file, "rb") as val_file:
            val_response = client.files.create(file=val_file, purpose="fine-tune")
        validation_file_id = val_response.id

    with console.status(f"Monitoring upload: {file_id} before finetuning...") as status:
        status.spinner_style = "dots"
        while True:
            file_status = get_file_status(file_id)
            if validation_file_id:
                validation_file_status = get_file_status(validation_file_id)

            if file_status == "processed" and (
                not validation_file_id or validation_file_status == "processed"
            ):
                console.log(f"[bold green]File {file_id} uploaded successfully!")
                if validation_file_id:
                    console.log(
                        f"[bold green]Validation file {validation_file_id} uploaded successfully!"
                    )
                break

            time.sleep(poll)

    additional_params: Dict[str, Union[str, Dict[str, Union[int, float, str]]]] = {}
    if hyperparameters_dict:
        additional_params["hyperparameters"] = hyperparameters_dict
    if validation_file:
        additional_params["validation_file"] = validation_file
    if model_suffix:
        additional_params["suffix"] = model_suffix

    job = client.fine_tuning.jobs.create(
        training_file=file_id,
        model=model,
        **additional_params,
    )
    if validation_file_id:
        console.log(
            f"[bold green]Fine-tuning job created with ID: {job.id} from file ID: {file_id} and validation_file ID: {validation_file_id}"
        )
    else:
        console.log(
            f"[bold green]Fine-tuning job created with ID: {job.id} from file ID: {file_id}"
        )
    watch(limit=5, poll=poll, screen=False)


@app.command(
    help="Cancel a fine-tuning job.",
)  # type: ignore[misc]
def cancel(
    id: str = typer.Argument(..., help="ID of the fine-tuning job to cancel"),
) -> None:
    with console.status(f"[bold red]Cancelling job {id}...", spinner="dots"):
        try:
            client.fine_tuning.jobs.cancel(id)
            console.log(f"[bold red]Job {id} cancelled successfully!")
        except Exception as e:
            console.log(f"[bold red]Error cancelling job {id}: {e}")


if __name__ == "__main__":
    app()
