"""Model reloader for hot-swapping fine-tuned models."""

from __future__ import annotations

import shutil
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from companion.config import Config


def reload_model(
    config: Config,
    new_model_path: Path,
    backup: bool = True,
) -> Path:
    """Reload the model with a new fine-tuned version.

    Args:
        config: Current configuration
        new_model_path: Path to new model directory or GGUF file
        backup: Whether to backup the old model

    Returns:
        Path to the active model
    """
    current_model = Path(config.model.inference.model_path).expanduser()

    # Validate new model exists
    if not new_model_path.exists():
        raise FileNotFoundError(f"New model not found: {new_model_path}")

    # Backup current model if it exists
    if backup and current_model.exists():
        backup_path = current_model.parent / f"{current_model.name}.backup"
        if backup_path.exists():
            shutil.rmtree(backup_path, ignore_errors=True)

        if current_model.is_dir():
            shutil.copytree(current_model, backup_path)
        else:
            shutil.copy2(current_model, backup_path)

        print(f"Backed up current model to: {backup_path}")

    # Copy new model to active location
    if current_model.exists():
        if current_model.is_dir():
            shutil.rmtree(current_model, ignore_errors=True)
        else:
            current_model.unlink()

    current_model.parent.mkdir(parents=True, exist_ok=True)

    if new_model_path.is_dir():
        shutil.copytree(new_model_path, current_model)
    else:
        shutil.copy2(new_model_path, current_model)

    print(f"Model reloaded: {new_model_path} -> {current_model}")

    return current_model


def get_model_status(config: Config) -> dict:
    """Get status of current model."""
    model_path = Path(config.model.inference.model_path).expanduser()

    status = {
        "path": str(model_path),
        "exists": model_path.exists(),
        "type": None,
        "size_mb": 0,
    }

    if model_path.exists():
        if model_path.is_dir():
            status["type"] = "directory"
            # Calculate directory size
            total_size = sum(
                f.stat().st_size for f in model_path.rglob("*") if f.is_file()
            )
            status["size_mb"] = round(total_size / (1024 * 1024), 2)
        else:
            status["type"] = "file"
            status["size_mb"] = round(model_path.stat().st_size / (1024 * 1024), 2)

    return status
