import torch
from typing import Dict, Any, Optional, List, Union

class PEFTTrainer:
    def __init__(self, model, peft: Any, lr: float = 0.0001, weight_decay: float = 0.0, device: Optional[str] = None, max_norm: float = 1.0):
        self.model = model
        self.max_norm = max_norm
        params = list(model.parameters())
        if params:
            self.device = device or str(params[0].device)
        else:
            self.device = device or 'cpu'
        if isinstance(self.device, str):
            self.device = torch.device(self.device)
        self.model.to(self.device)

        if isinstance(peft, list):
            self.pefts = peft
        else:
            self.pefts = [peft]

        self.peft_params = []
        for p in self.pefts:
            self.peft_params.extend(list(p.parameters()))

        if not self.peft_params:
            raise ValueError("No trainable parameters found in PEFT config")
        self.optimizer = torch.optim.AdamW(self.peft_params, lr=lr, weight_decay=weight_decay)

    def train_step(self, batch: Dict[str, Any]):
        self.model.train()
        batch = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
        out = self.model(**batch)
        if hasattr(out, 'loss') and out.loss is not None:
            loss = out.loss
        elif isinstance(out, dict) and 'loss' in out:
            loss = out['loss']
        elif isinstance(out, tuple) and len(out) > 0:
            loss = out[0]
        else:
            raise ValueError('no loss in model output')
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.peft_params, self.max_norm)
        self.optimizer.step()
        self.optimizer.zero_grad()
        return loss.item()