import os
import functools
import torch
import torch.nn as nn
from typing import Optional, Dict, Any
import wandb
import numpy as np

__all__ = ["handle_keyboard_interrupt", "calculate_metrics", "WandBHandler"]

def handle_keyboard_interrupt(func):

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except KeyboardInterrupt:
            print("\n[INFO] KeyboardInterrupt caught — finishing WandB run and exiting.")
            wandb.finish()
    return wrapper

def calculate_metrics(cm) -> Dict[str, float]:
    tn = cm[0][0]
    fp = cm[0][1]
    fn = cm[1][0]
    tp = cm[1][1]

    recall    = tp / max(tp + fn, 1e-10)
    precision = tp / max(tp + fp, 1e-10)
    far       = fp / max(fp + tn, 1e-10)   # real predicted as fake
    frr       = fn / max(fn + tp, 1e-10)   # fake predicted as real
    f1        = 2 * recall * precision / max(recall + precision, 1e-10)
    accuracy  = (tp + tn) / max(tp + tn + fp + fn, 1e-10)

    return {
        "recall (fake detection rate)":                     recall,
        "precision (of all predicted fake, actually fake)": precision,
        "false alarm rate (real predicted as fake)":        far,
        "false rejection rate (fake predicted as real)":    frr,
        "f1 score":                                         f1,
        "accuracy":                                         accuracy,
    }


class WandBHandler:

    def __init__(
        self,
        project_name: str,
        run_name: Optional[str],
        config: Dict[str, Any],
        entity: Optional[str] = "axoryai-axory",
        resume_run: Optional[str] = None,
    ):
        self.project_name = project_name
        self.run_name     = run_name
        self.config       = config
        self.entity       = entity
        self.resume_run   = resume_run
        self.run          = None
        self._initialize()

    def _initialize(self):
        if self.run is not None:
            return

        init_kwargs = dict(
            project=self.project_name,
            name=self.run_name,
            config=self.config,
        )
        if self.entity:
            init_kwargs["entity"] = self.entity

        if self.resume_run:
            print(f"[WandB] Resuming run id={self.resume_run}")
            init_kwargs.update(id=self.resume_run, resume="allow")

        self.run = wandb.init(**init_kwargs)

    def log_metrics(self, metrics: Dict[str, Any]):
        """Log a dict of scalar or image metrics to WandB."""
        self._initialize()
        wandb.log(metrics)

    def save_checkpoint(
        self,
        model: nn.Module,
        filename: str,
        epoch: int,
        best_metric: float,
        metadata: Optional[Dict] = None,
        is_best: bool = False,
        extra_metrics: Optional[Dict] = None,
    ):
        meta = {
            "epoch":       epoch,
            "is_best":     is_best,
            "best_metric": best_metric,
            **(metadata or {}),
            **(extra_metrics or {}),
        }

        torch.save(model.state_dict(), filename)

        artifact = wandb.Artifact(
            name=filename.replace(".pth", "").replace("/", "-"),
            type="model",
            metadata=meta,
        )
        artifact.add_file(filename)
        wandb.log_artifact(artifact)

        if not is_best and os.path.exists(filename):
            os.remove(filename)

    def load_model(self, artifact_name: str) -> Dict:
        self._initialize()
        artifact     = wandb.use_artifact(artifact_name, type="model")
        artifact_dir = artifact.download()
        print(f"[WandB] Artifact downloaded to: {artifact_dir}")

        pth_files = [f for f in os.listdir(artifact_dir) if f.endswith(".pth")]
        if not pth_files:
            raise FileNotFoundError("No .pth file found in the artifact directory.")

        state_dict = torch.load(
            os.path.join(artifact_dir, pth_files[0]), map_location="cpu"
        )
        return {"model_state_dict": state_dict, **artifact.metadata}

    def finish_run(self):
        if self.run:
            self.run.finish()
            self.run = None