"""Weights & Biases logging wrapper."""

from dataclasses import asdict
from datetime import datetime
from pathlib import Path
from typing import Any, Optional

from ..config import Config, WandbConfig


class WandbLogger:
    """Weights & Biases logging wrapper with graceful fallback."""

    def __init__(self, wandb_config: WandbConfig, run_config: Config):
        """Initialize Wandb logger.

        Args:
            wandb_config: Wandb configuration.
            run_config: Full run configuration for logging.
        """
        self.enabled = wandb_config.enabled
        self.wandb = None
        self.run = None

        if self.enabled:
            try:
                import wandb

                self.wandb = wandb

                run_name = wandb_config.run_name or (
                    f"{run_config.model_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
                )

                tags = list(wandb_config.tags)
                language = run_config.data.language
                if language and language not in tags:
                    tags.append(language)

                self.run = wandb.init(
                    project=wandb_config.project,
                    entity=wandb_config.entity,
                    name=run_name,
                    tags=tags,
                    config=self._config_to_dict(run_config),
                )

                # Define metric axes
                wandb.define_metric("global_step")
                wandb.define_metric("epoch")
                wandb.define_metric("train/*", step_metric="global_step")
                wandb.define_metric("dev_eval/*", step_metric="epoch")
                wandb.define_metric("test_eval/*", step_metric="epoch")
                wandb.define_metric("train_eval/*", step_metric="epoch")
                wandb.define_metric("dev_fulltext/*", step_metric="epoch")
                wandb.define_metric("test_fulltext/*", step_metric="epoch")

                print(f"Wandb initialized: {wandb_config.project}/{run_name}")
            except ImportError:
                print("Warning: wandb not installed, logging disabled")
                self.enabled = False
            except Exception as e:
                print(f"Warning: Failed to initialize wandb: {e}")
                self.enabled = False

    def _config_to_dict(self, config: Config) -> dict:
        """Convert config to dict for wandb logging."""
        return {
            "data": asdict(config.data),
            "model": {
                k: v
                for k, v in asdict(config.model).items()
                if not k.startswith("num_")  # Exclude dynamically set values
            },
            "training": asdict(config.training),
            "seed": config.seed,
            "device": config.device,
        }

    def log(self, metrics: dict[str, Any], step: Optional[int] = None) -> None:
        """Log metrics.

        Args:
            metrics: Dictionary of metric name to value.
            step: Optional step number.
        """
        if self.enabled and self.wandb:
            self.wandb.log(metrics, step=step)

    def log_table(
        self, name: str, columns: list[str], data: list[list[Any]]
    ) -> None:
        """Log a table.

        Args:
            name: Table name.
            columns: Column names.
            data: Table data (list of rows).
        """
        if self.enabled and self.wandb:
            table = self.wandb.Table(columns=columns, data=data)
            self.wandb.log({name: table})

    def log_artifact(
        self, path: Path, name: str, artifact_type: str = "model"
    ) -> None:
        """Log model or data artifact.

        Args:
            path: Path to artifact file.
            name: Artifact name.
            artifact_type: Type of artifact (e.g., "model", "dataset").
        """
        if self.enabled and self.wandb:
            artifact = self.wandb.Artifact(name, type=artifact_type)
            if path.is_dir():
                artifact.add_dir(str(path))
            else:
                artifact.add_file(str(path))
            self.run.log_artifact(artifact)

    def watch(self, model: Any, log: str = "gradients", log_freq: int = 100) -> None:
        """Watch model gradients and parameters.

        Args:
            model: PyTorch model.
            log: What to log ("gradients", "parameters", "all").
            log_freq: Logging frequency.
        """
        if self.enabled and self.wandb:
            self.wandb.watch(model, log=log, log_freq=log_freq)

    def summary(self, metrics: dict[str, Any]) -> None:
        """Update run summary with final metrics.

        Args:
            metrics: Dictionary of summary metrics.
        """
        if self.enabled and self.run:
            for key, value in metrics.items():
                self.run.summary[key] = value

    def finish(self) -> None:
        """Finish the run."""
        if self.enabled and self.wandb:
            self.wandb.finish()

    @property
    def run_id(self) -> Optional[str]:
        """Get the run ID."""
        if self.run:
            return self.run.id
        return None

    @property
    def run_url(self) -> Optional[str]:
        """Get the run URL."""
        if self.run:
            return self.run.get_url()
        return None
