#!/usr/bin/env python3
"""
Publish script for Databricks MCP Server
Builds and publishes the package to PyPI
"""

import os
import shutil
import subprocess
import sys
from pathlib import Path


def run_command(cmd: list, cwd: Path = None) -> bool:
    """Run a command and return success status."""
    print(f"\n{'='*60}")
    print(f"Running: {' '.join(cmd)}")
    print(f"{'='*60}\n")
    
    try:
        result = subprocess.run(
            cmd,
            cwd=cwd,
            check=True,
            capture_output=False
        )
        return True
    except subprocess.CalledProcessError as e:
        print(f"Command failed with exit code {e.returncode}")
        return False


def clean_build_artifacts(project_dir: Path) -> None:
    """Clean up build artifacts."""
    print("\nCleaning build artifacts...")
    
    dirs_to_remove = [
        project_dir / "dist",
        project_dir / "build",
        project_dir / "src" / "mareana_databricks_mcp_server.egg-info",
        project_dir / "src" / "databricks_mcp_server.egg-info",
    ]
    
    for d in dirs_to_remove:
        if d.exists():
            print(f"  Removing {d}")
            shutil.rmtree(d)
    
    print("Clean complete.")


def build_package(project_dir: Path) -> bool:
    """Build the package."""
    print("\n" + "="*60)
    print("Building package...")
    print("="*60)
    
    return run_command([sys.executable, "-m", "build"], cwd=project_dir)


def publish_to_pypi(project_dir: Path, test: bool = False) -> bool:
    """Publish to PyPI or TestPyPI."""
    if test:
        print("\n" + "="*60)
        print("Publishing to TestPyPI...")
        print("="*60)
        return run_command([
            sys.executable, "-m", "twine", "upload",
            "--repository-url", "https://test.pypi.org/legacy/",
            "dist/*"
        ], cwd=project_dir)
    else:
        print("\n" + "="*60)
        print("Publishing to PyPI...")
        print("="*60)
        return run_command([
            sys.executable, "-m", "twine", "upload",
            "dist/*"
        ], cwd=project_dir)


def main():
    """Main entry point."""
    project_dir = Path(__file__).parent
    
    print("="*60)
    print("Databricks MCP Server - Build & Publish")
    print("="*60)
    
    # Check for required tools
    print("\nChecking requirements...")
    
    try:
        subprocess.run([sys.executable, "-m", "build", "--version"], 
                       capture_output=True, check=True)
    except (subprocess.CalledProcessError, FileNotFoundError):
        print("Installing build package...")
        # Use uv pip if available, fallback to pip
        try:
            subprocess.run(["uv", "pip", "install", "build"], check=True)
        except FileNotFoundError:
            subprocess.run([sys.executable, "-m", "pip", "install", "build"], check=True)
    
    try:
        subprocess.run([sys.executable, "-m", "twine", "--version"], 
                       capture_output=True, check=True)
    except (subprocess.CalledProcessError, FileNotFoundError):
        print("Installing twine package...")
        # Use uv pip if available, fallback to pip
        try:
            subprocess.run(["uv", "pip", "install", "twine"], check=True)
        except FileNotFoundError:
            subprocess.run([sys.executable, "-m", "pip", "install", "twine"], check=True)
    
    # Parse arguments
    test_mode = "--test" in sys.argv
    skip_build = "--skip-build" in sys.argv
    clean_only = "--clean" in sys.argv
    
    # Clean
    clean_build_artifacts(project_dir)
    
    if clean_only:
        print("\nClean only mode - done.")
        return
    
    # Build
    if not skip_build:
        if not build_package(project_dir):
            print("\n❌ Build failed!")
            sys.exit(1)
        print("\n✅ Build successful!")
    
    # Publish
    if "--no-publish" not in sys.argv:
        if test_mode:
            print("\n⚠️  Publishing to TestPyPI...")
        else:
            print("\n⚠️  Publishing to PyPI (production)...")
        
        # Check for PyPI token
        if not os.getenv("TWINE_PASSWORD") and not os.getenv("PYPI_TOKEN"):
            print("\n⚠️  No PyPI token found in environment.")
            print("Set TWINE_PASSWORD or PYPI_TOKEN environment variable.")
            
            token = input("\nEnter PyPI API token (or press Enter to skip publish): ").strip()
            if not token:
                print("Skipping publish.")
                return
            os.environ["TWINE_USERNAME"] = "__token__"
            os.environ["TWINE_PASSWORD"] = token
        
        if os.getenv("PYPI_TOKEN"):
            os.environ["TWINE_USERNAME"] = "__token__"
            os.environ["TWINE_PASSWORD"] = os.getenv("PYPI_TOKEN")
        
        if not publish_to_pypi(project_dir, test=test_mode):
            print("\n❌ Publish failed!")
            sys.exit(1)
        
        print("\n✅ Published successfully!")
        if test_mode:
            print("   Package available at: https://test.pypi.org/project/mareana-databricks-mcp-server/")
        else:
            print("   Package available at: https://pypi.org/project/mareana-databricks-mcp-server/")
    else:
        print("\n✅ Build complete (publish skipped).")


if __name__ == "__main__":
    main()
