# SPDX-License-Identifier: AGPL-3.0-or-later
# Copyright (C) 2026 Assem Sabry

"""Shade: Fully automatic censorship removal for language models."""

from importlib.metadata import version, PackageNotFoundError

try:
    __version__ = version("shadeai")
except PackageNotFoundError:
    __version__ = "2.2.1"

from .config import Settings, QuantizationMethod, RowNormalization, DatasetSpecification
from .model import Model, AbliterationParameters
from .cli import run_doctor
from .main import main as run_optimization
from .server import start_server

# Useful wrappers for the Shade framework
def login(token: str):
    """Login to Hugging Face Hub."""
    from huggingface_hub import login as hf_login
    hf_login(token=token, add_to_git_credential=True)

def download(model_id: str, allow_patterns: str = None):
    """Download a model from Hugging Face."""
    from huggingface_hub import snapshot_download
    snapshot_download(repo_id=model_id, allow_patterns=allow_patterns)

def prune(clean_all: bool = False):
    """Clean up temporary files and cache to save space."""
    from .cli import clear_cmd
    clear_cmd.callback(all=clean_all)

def export_to_ollama(model_path: str, name: str = None):
    """Integrate a saved model with Ollama."""
    from .cli import ollama
    ollama.callback(model_path=model_path, name=name)

def export_to_gguf(model_path: str, quant_type: str = "q4_k_m", output_path: str = None):
    """Export a HuggingFace model to GGUF format."""
    from .gguf_export import convert_to_gguf
    return convert_to_gguf(model_path=model_path, output_path=output_path, quant_type=quant_type)

def benchmark(model_id: str):
    """Run a quality benchmark on a model."""
    from .cli import benchmark as cli_benchmark
    cli_benchmark.callback(model_id=model_id)

def serve(model_id: str = None, host: str = "127.0.0.1", port: int = 8000):
    """Start local web server."""
    from .cli import serve as cli_serve
    cli_serve.callback(host=host, port=port, model_id=model_id)

def install_cuda(url: bool = False):
    """Download and install CUDA helper."""
    from .cli import run_install_cuda
    run_install_cuda(url=url)

__all__ = [
    "__version__",
    "Settings",
    "QuantizationMethod",
    "RowNormalization",
    "DatasetSpecification",
    "Model",
    "AbliterationParameters",
    "run_doctor",
    "run_optimization",
    "start_server",
    "login",
    "download",
    "prune",
    "export_to_ollama",
    "export_to_gguf",
    "benchmark",
    "serve",
    "install_cuda",
]
