import torch
import logging
from pathlib import Path
from app.models.unet_model import BuildUNet
from app.core.config import settings
from app.core.exceptions import ModelLoadError

logger = logging.getLogger(__name__)

class ModelLoader:
    def __init__(self):
        self.model = None
        self.device = None
        self._load_device()
    
    def _load_device(self):
        if settings.MODEL_DEVICE == "cuda" and torch.cuda.is_available():
            self.device = torch.device("cuda")
            logger.info(f"Using GPU: {torch.cuda.get_device_name(0)}")
        else:
            self.device = torch.device("cpu")
            logger.info("Using CPU")
    
    def load_model(self):
        try:
            model_path = Path(settings.MODEL_PATH)
            if not model_path.exists():
                raise ModelLoadError(f"Model file not found at {settings.MODEL_PATH}")
            
            logger.info(f"Loading model from {settings.MODEL_PATH}")
            self.model = BuildUNet(num_classes=settings.NUM_CLASSES)
            
            checkpoint = torch.load(model_path, map_location=self.device)
            if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
                self.model.load_state_dict(checkpoint["model_state_dict"])
            else:
                self.model.load_state_dict(checkpoint)
            
            self.model.to(self.device)
            self.model.eval()
            logger.info("Model loaded successfully")
            
        except Exception as e:
            logger.error(f"Failed to load model: {str(e)}")
            raise ModelLoadError(f"Failed to load model: {str(e)}")
    
    def get_model(self):
        if self.model is None:
            raise ModelLoadError("Model not loaded. Call load_model() first.")
        return self.model
    
    def get_device(self):
        return self.device

model_loader = ModelLoader()