"""Logging to WandB"""

import uuid
from pathlib import Path
from typing import Union

from wandb import Api, Image, Run, Table
from wandb import finish as wandb_finish
from wandb import init as wandb_init

from .logger import Logger


class WBLogger(Logger):
    """The logger for wandb, implmenting the Logger interface"""

    _run: Run
    _record_tables: dict[str, Table]

    def __init__(self, project_name: str = "ml_base", run_name: str = str(uuid.uuid4())):
        """A logger which saves to WandB

        Args:
            project_name (str, optional): The project name in wandb console
        """
        super().__init__()
        # Current run
        self._run = wandb_init(
            project=project_name,
            name=run_name,
            resume="allow",
            id=f"{project_name}_{run_name}",
        )
        self._record_tables = {}

    def log_record(self, record: dict[str, Union[str, float]], table_name: str = "main"):
        """Store record in a table incrementally

        Args:
            record (dict[str,float]): Mapping from column name to value to record
            table_name (str): The table in which to store the record
        """
        # Make table if it does not exist
        if self._record_tables.get(table_name) is None:
            colums = list(record.keys())
            self._record_tables[table_name] = Table(columns=colums, log_mode="INCREMENTAL")
        # Collect values in correct order
        values: list[Union[Image, str, float]] = []
        for col in self._record_tables[table_name].columns:
            record_val = record[col]
            if isinstance(record_val, str) and Path(record_val).exists():
                file_path = Path(record_val)
                if file_path.suffix.lower() in ["png", "jpg", "jpeg"]:
                    values.append(Image(record_val))
                else:
                    print(f"unexpected file type: {file_path.suffix.lower()}")
                    # Keep value
                    values.append(record_val)
            else:
                # float or non path str, append raw value
                values.append(record_val)
        self._record_tables[table_name].add_data(values)
        # Upload scalars to wandb table
        self._run.log(data={table_name: self._record_tables[table_name]})

    def log_text(self, string: str, table_name: str = "logs"):
        """Store text in a table incrementally

        Args:
            text (str): Text to save in table
            table_name (str): The table in which to store the record
        """
        # Make table if it does not exist\
        if self._record_tables.get(table_name) is None:
            self._record_tables[table_name] = Table(columns=["text"], log_mode="INCREMENTAL")
        self._record_tables[table_name].add_data((string,))
        # Upload text
        self._run.log(data={table_name: self._record_tables[table_name]})

    def __del__(self):
        """Disconnect from wandb when logger is deleted"""
        wandb_finish()
        # Flush local files
        Api().flush()
