#!/usr/bin/env python3
"""Automated notebook execution, PDF conversion, and Slack notification."""

import logging
import os
import subprocess
import sys
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional

import requests
import yaml


def setup_logging(log_dir: Path) -> logging.Logger:
    """Set up logging with file and console handlers."""
    log_dir.mkdir(parents=True, exist_ok=True)
    log_file = log_dir / f"notebook_automation_{datetime.now().strftime('%Y%m%d')}.log"

    logger = logging.getLogger("notebook_automation")
    logger.setLevel(logging.INFO)

    # File handler
    file_handler = logging.FileHandler(log_file)
    file_handler.setLevel(logging.INFO)

    # Console handler
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)

    # Formatter
    formatter = logging.Formatter(
        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    )
    file_handler.setFormatter(formatter)
    console_handler.setFormatter(formatter)

    logger.addHandler(file_handler)
    logger.addHandler(console_handler)

    return logger


def load_config(config_path: Path) -> Dict:
    """Load configuration from YAML file."""
    with open(config_path, "r") as f:
        return yaml.safe_load(f)


def execute_notebook(notebook_path: Path, logger: logging.Logger) -> bool:
    """Execute notebook in-place."""
    try:
        logger.info(f"Starting notebook execution: {notebook_path}")

        # Execute notebook without PDF conversion
        cmd = [
            "jupyter",
            "nbconvert",
            "--execute",
            "--to",
            "notebook",
            "--inplace",
            "--ExecutePreprocessor.timeout=3600",  # 1 hour timeout
            "--allow-errors",  # Continue execution even if there are errors
            str(notebook_path),
        ]

        # Set environment to disable tqdm progress bars
        env = os.environ.copy()
        env["TQDM_DISABLE"] = "1"

        result = subprocess.run(
            cmd, capture_output=True, text=True, check=False, env=env
        )

        if result.returncode != 0:
            logger.error(f"Notebook execution failed: {result.stderr}")
            return False

        logger.info("Notebook executed successfully")
        return True

    except Exception as e:
        logger.error(f"Error during notebook execution: {e}")
        return False


def extract_plots_from_notebook(
    notebook_path: Path, reports_dir: Path, logger: logging.Logger
) -> list[Path]:
    """Extract plots from executed notebook and save as JPEGs."""
    import json
    import base64
    from PIL import Image
    import io

    try:
        logger.info(f"Extracting plots from notebook: {notebook_path}")

        # Read the executed notebook
        with open(notebook_path, "r") as f:
            nb = json.load(f)

        # Define the cells we want to extract plots from
        target_cells = [
            "get_preference_counts(interesting_clips)",
            "get_preference_counts(user_intersting_clips_3p5)",
        ]

        extracted_plots = []
        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M")

        # Find and extract plots from target cells
        for cell in nb["cells"]:
            if cell["cell_type"] == "code":
                source = "".join(cell.get("source", []))

                for idx, target in enumerate(target_cells):
                    if target in source:
                        # Look for image in outputs
                        for output in cell.get("outputs", []):
                            if output.get("output_type") == "display_data":
                                data = output.get("data", {})
                                if "image/png" in data:
                                    plot_data = data["image/png"]
                                    logger.info(f"Found plot for: {target}")

                                    # Decode and save the image
                                    plot_name = (
                                        "interesting_clips"
                                        if idx == 0
                                        else "user_intersting_clips_3p5"
                                    )
                                    output_path = (
                                        reports_dir
                                        / f"preference_plot_{plot_name}_{timestamp}.jpg"
                                    )

                                    # Decode base64 and save as JPEG
                                    image_data = base64.b64decode(plot_data)

                                    # Convert PNG to JPEG using PIL
                                    img = Image.open(io.BytesIO(image_data))
                                    # Convert RGBA to RGB if necessary
                                    if img.mode in ("RGBA", "LA", "P"):
                                        background = Image.new(
                                            "RGB", img.size, (255, 255, 255)
                                        )
                                        if img.mode == "P":
                                            img = img.convert("RGBA")
                                        background.paste(
                                            img,
                                            mask=img.split()[-1]
                                            if img.mode == "RGBA"
                                            else None,
                                        )
                                        img = background

                                    img.save(output_path, "JPEG", quality=95)
                                    logger.info(f"Plot saved to: {output_path}")
                                    extracted_plots.append(output_path)
                                    break

        if not extracted_plots:
            logger.error("No plots found in notebook outputs")
            return []

        logger.info(f"Successfully extracted {len(extracted_plots)} plot(s)")
        return extracted_plots

    except Exception as e:
        logger.error(f"Error extracting plots from notebook: {e}")
        return []


def send_slack_webhook(webhook_url: str, message: str, logger: logging.Logger) -> bool:
    """Send a message to Slack via webhook."""
    try:
        payload = {"text": message}
        response = requests.post(webhook_url, json=payload, timeout=30)
        response.raise_for_status()
        logger.info("Slack message sent successfully via webhook")
        return True
    except Exception as e:
        logger.error(f"Failed to send Slack webhook: {e}")
        return False


def get_channel_id(
    bot_token: str, channel_name: str, logger: logging.Logger
) -> Optional[str]:
    """
    Get Slack channel ID from channel name.

    Args:
        bot_token: Slack bot token
        channel_name: Channel name (with or without #)
        logger: Logger instance

    Returns:
        Channel ID if found, None otherwise
    """
    try:
        # Remove # if present
        clean_name = channel_name.lstrip("#")

        headers = {"Authorization": f"Bearer {bot_token}"}
        response = requests.get(
            "https://slack.com/api/conversations.list",
            headers=headers,
            params={"types": "public_channel,private_channel", "limit": 1000},
            timeout=30,
        )

        result = response.json()
        if result.get("ok"):
            for channel in result.get("channels", []):
                if channel.get("name") == clean_name:
                    return channel.get("id")

        logger.error(f"Channel '{channel_name}' not found")
        return None
    except Exception as e:
        logger.error(f"Failed to get channel ID: {e}")
        return None


def upload_file_to_slack_webhook(
    webhook_url: str,
    bot_token: Optional[str],
    channel: str,
    file_path: Path,
    message: str,
    logger: logging.Logger,
) -> bool:
    """
    Upload a file to Slack channel.

    Args:
        webhook_url: Slack webhook URL (for text message)
        bot_token: Slack bot token for file upload
        channel: Slack channel name or ID
        file_path: Path to file to upload
        message: Message to accompany the file
        logger: Logger instance

    Returns:
        True if successful, False otherwise
    """
    try:
        # Use bot token if available for file upload (better for files)
        if bot_token:
            headers = {"Authorization": f"Bearer {bot_token}"}

            # Check if channel is already an ID (starts with C, G, D) or a name
            if channel.startswith(("C", "G", "D")):
                # Already a channel ID
                channel_id = channel
                logger.info(f"Using provided channel ID: {channel_id}")
            else:
                # Get channel ID from channel name
                channel_id = get_channel_id(bot_token, channel, logger)
                if not channel_id:
                    logger.error(f"Could not find channel ID for '{channel}'")
                    send_slack_webhook(webhook_url, message, logger)
                    return False

            # Step 1: Get upload URL
            file_size = file_path.stat().st_size

            upload_url_response = requests.post(
                "https://slack.com/api/files.getUploadURLExternal",
                headers=headers,
                data={
                    "filename": file_path.name,
                    "length": file_size,
                },
                timeout=30,
            )

            upload_data = upload_url_response.json()
            if not upload_data.get("ok"):
                error = upload_data.get("error")
                logger.error(f"Failed to get upload URL: {error}")
                send_slack_webhook(webhook_url, message, logger)
                return False

            upload_url = upload_data["upload_url"]
            file_id = upload_data["file_id"]

            # Step 2: Upload file to the URL
            with open(file_path, "rb") as f:
                upload_response = requests.post(
                    upload_url,
                    files={"file": f},
                    timeout=120,
                )

            if upload_response.status_code != 200:
                logger.error(
                    f"File upload to URL failed: {upload_response.status_code}"
                )
                send_slack_webhook(webhook_url, message, logger)
                return False

            # Step 3: Complete the upload and share to channel
            import json as json_module

            complete_response = requests.post(
                "https://slack.com/api/files.completeUploadExternal",
                headers=headers,
                data={
                    "files": json_module.dumps(
                        [{"id": file_id, "title": file_path.name}]
                    ),
                    "channel_id": channel_id,
                    "initial_comment": message,
                },
                timeout=30,
            )

            result = complete_response.json()
            if result.get("ok"):
                logger.info(f"File uploaded to Slack successfully: {file_path.name}")
                return True
            else:
                error = result.get("error")
                logger.error(f"Slack file upload failed: {error}")
                logger.error(f"Full response: {result}")
                # Fall back to just sending a message
                send_slack_webhook(webhook_url, message, logger)
                return False
        else:
            # Without bot token, just send a message via webhook
            logger.warning("No bot token provided, sending message only (no file)")
            return send_slack_webhook(webhook_url, message, logger)

    except Exception as e:
        logger.error(f"Failed to upload file to Slack: {e}")
        # Try to send error message
        send_slack_webhook(
            webhook_url, f"{message}\n\n⚠️ Failed to upload file: {e}", logger
        )
        return False


def cleanup_old_files(
    directory: Path, days_to_keep: int, logger: logging.Logger
) -> None:
    """
    Remove files older than specified days.

    Args:
        directory: Directory to clean
        days_to_keep: Number of days to keep files
        logger: Logger instance
    """
    try:
        cutoff_time = datetime.now().timestamp() - (days_to_keep * 86400)
        removed_count = 0

        for file_path in directory.glob("*"):
            if file_path.is_file() and file_path.stat().st_mtime < cutoff_time:
                file_path.unlink()
                removed_count += 1

        if removed_count > 0:
            logger.info(f"Cleaned up {removed_count} old files from {directory}")

    except Exception as e:
        logger.warning(f"Failed to cleanup old files: {e}")


def main() -> int:
    """
    Main execution function.

    Returns:
        Exit code (0 for success, 1 for failure)
    """
    import argparse

    # Parse command line arguments
    parser = argparse.ArgumentParser(
        description="Execute notebook and send plot to Slack"
    )
    parser.add_argument(
        "--skip-notebook",
        action="store_true",
        help="Skip notebook execution and upload most recent plot",
    )
    args = parser.parse_args()

    # Determine script directory and config path
    script_dir = Path(__file__).parent
    config_path = script_dir / "config.yaml"

    # Setup logging first
    log_dir = script_dir / "logs"
    logger = setup_logging(log_dir)

    try:
        # Load configuration
        logger.info("Loading configuration")
        config = load_config(config_path)

        # Setup paths
        notebook_path = script_dir / config["notebook_path"]
        reports_dir = script_dir / config["reports_dir"]
        reports_dir.mkdir(parents=True, exist_ok=True)

        # Generate timestamp
        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M")

        # Execute notebook if not skipped
        if not args.skip_notebook:
            logger.info("Executing notebook...")
            if not execute_notebook(notebook_path, logger):
                error_msg = f"❌ Notebook execution failed at {timestamp}\nCheck logs for details."
                send_slack_webhook(config["slack_webhook_url"], error_msg, logger)
                return 1
        else:
            logger.info("Skipping notebook execution")

        # Extract plots from notebook
        output_jpegs = extract_plots_from_notebook(notebook_path, reports_dir, logger)

        if not output_jpegs:
            # If extraction fails, try to find existing plots
            logger.warning("Plot extraction failed, looking for existing plots...")
            jpeg_files = sorted(
                reports_dir.glob("preference_plot_*.jpg"),
                key=lambda p: p.stat().st_mtime,
                reverse=True,
            )
            if not jpeg_files:
                error_msg = "❌ No plots found in notebook output or reports directory\nCheck logs for details."
                send_slack_webhook(config["slack_webhook_url"], error_msg, logger)
                return 1
            # Use the two most recent plots
            output_jpegs = jpeg_files[:2]
            logger.info(f"Using {len(output_jpegs)} existing plot(s)")

        # Send plots to Slack
        for idx, output_jpeg in enumerate(output_jpegs, 1):
            plot_type = (
                "interesting_clips"
                if "interesting_clips" in output_jpeg.name
                else "user_intersting_clips_3p5"
            )
            success_msg = (
                f"✅ Preference Plot Generated ({idx}/{len(output_jpegs)})\n"
                f"Dataset: {plot_type}\n"
                f"Timestamp: {timestamp}\n"
                f"File: {output_jpeg.name}"
            )

            upload_success = upload_file_to_slack_webhook(
                webhook_url=config["slack_webhook_url"],
                bot_token=config.get("slack_bot_token"),
                channel=config["slack_channel"],
                file_path=output_jpeg,
                message=success_msg,
                logger=logger,
            )

            if not upload_success:
                logger.warning(
                    f"Upload of {output_jpeg.name} had issues, continuing..."
                )

        # Cleanup old files (keep last 1 day for reports, 7 days for logs)
        cleanup_old_files(reports_dir, 1, logger)
        cleanup_old_files(log_dir, 7, logger)

        logger.info("Automation completed successfully")
        return 0

    except Exception as e:
        logger.error(f"Fatal error in main execution: {e}", exc_info=True)
        try:
            config = load_config(config_path)
            error_msg = f"❌ Critical error in notebook automation:\n{str(e)}"
            send_slack_webhook(config["slack_webhook_url"], error_msg, logger)
        except Exception:
            pass
        return 1


if __name__ == "__main__":
    sys.exit(main())