"""CLI tool for ceil-dlp"""

from pathlib import Path

import typer
import yaml
from rich.console import Console
from rich.panel import Panel
from rich.text import Text

console = Console()

app = typer.Typer(help="ceil-dlp: Open-Source DLP for LLMs and Agentic Workflows")


@app.callback(invoke_without_command=True)
def main_callback(ctx: typer.Context) -> None:
    """ceil-dlp: Open-Source DLP for LLMs and Agentic Workflows"""
    if ctx.invoked_subcommand is None:
        typer.echo(ctx.get_help())


CALLBACK_WRAPPER_TEMPLATE = '''"""Local callback wrapper for ceil-dlp.

This file is auto-generated by 'ceil-dlp install'. It provides a local import
path for LiteLLM to use ceil-dlp as a callback.
"""

import os
import sys
from pathlib import Path

# Add this directory to sys.path so LiteLLM can import this module
# regardless of where it's run from
_config_dir = Path(__file__).parent
if str(_config_dir) not in sys.path:
    sys.path.insert(0, str(_config_dir))

# Import from the installed ceil-dlp package
# This import is needed for LiteLLM to discover the handler instance
from ceil_dlp.ceil_dlp_callback import (  # noqa: E402
    proxy_handler_instance,  # noqa: F401  # type: ignore[unused-import]
)

# Set CEIL_DLP_CONFIG_PATH to use ceil-dlp.yaml in the same directory as this file
_config_file = _config_dir / "ceil-dlp.yaml"

if _config_file.exists():
    os.environ["CEIL_DLP_CONFIG_PATH"] = str(_config_file)
'''

DEFAULT_CEIL_DLP_CONFIG = """# ceil-dlp Configuration
# This file was auto-generated by 'ceil-dlp install'
# Customize as needed for your environment

mode: enforce  # observe | warn | enforce

policies:
  # High-risk items - block by default
  credit_card:
    action: block
    enabled: true
  ssn:
    action: block
    enabled: true
  api_key:
    action: block
    enabled: true
  pem_key:
    action: block
    enabled: true
  jwt_token:
    action: block
    enabled: true
  high_entropy_token:
    action: block
    enabled: true

  # Medium-risk items - mask by default
  email:
    action: mask
    enabled: true
  phone:
    action: mask
    enabled: true

# Audit logging (optional)
# audit_log_path: /var/log/ceil-dlp/audit.log

# PII types to detect (if not specified, all are enabled)
# enabled_pii_types:
#   - credit_card
#   - ssn
#   - email
#   - phone
#   - api_key
"""


@app.command()
def install(
    litellm_config: Path = typer.Argument(
        ...,
        help="Path to LiteLLM config.yaml file",
        exists=True,
        file_okay=True,
        dir_okay=False,
    ),
    update_config: bool = typer.Option(
        True,
        "--update-config/--no-update-config",
        help="Automatically update LiteLLM config.yaml to include ceil-dlp callback",
    ),
) -> None:
    """Install ceil-dlp in the same directory as LiteLLM config file.

    This command:
    1. Creates a local ceil_dlp_callback.py wrapper in the config directory
    2. Creates a starter ceil-dlp.yaml if it doesn't exist
    3. Optionally updates LiteLLM config.yaml to include the callback
    """
    config_dir = litellm_config.parent
    callback_file = config_dir / "ceil_dlp_callback.py"
    ceil_dlp_config = config_dir / "ceil-dlp.yaml"

    console.print("\n[bold blue]installing ceil-dlp[/bold blue]")
    console.print(f"config directory: [cyan]{config_dir}[/cyan]")
    console.print(f"litellm config: [cyan]{litellm_config}[/cyan]\n")

    # create callback wrapper
    if callback_file.exists():
        console.print(f"[dim]{callback_file.name}[/dim] already exists, skipping...")
    else:
        callback_file.write_text(CALLBACK_WRAPPER_TEMPLATE)
        console.print(f"created [cyan]{callback_file.name}[/cyan]")

    # create ceil-dlp.yaml if it doesn't exist
    if ceil_dlp_config.exists():
        console.print(f"[dim]{ceil_dlp_config.name}[/dim] already exists, skipping...")
    else:
        ceil_dlp_config.write_text(DEFAULT_CEIL_DLP_CONFIG)
        console.print(f"created [cyan]{ceil_dlp_config.name}[/cyan]")

    # update LiteLLM config if requested
    if update_config:
        # calculate callback path
        # The wrapper adds its directory to sys.path, so we can always use
        # the simple module name regardless of where LiteLLM is run from
        callback_path = "ceil_dlp_callback.proxy_handler_instance"

        try:
            with litellm_config.open() as f:
                config = yaml.safe_load(f) or {}

            # check if callback is already configured
            callbacks = config.get("litellm_settings", {}).get("callbacks", [])

            if isinstance(callbacks, str):
                callbacks = [callbacks]

            # check if already configured (handle both string and list formats)
            already_configured = False
            if isinstance(callbacks, list):
                already_configured = any(
                    callback_path in str(cb) or "ceil_dlp_callback" in str(cb) for cb in callbacks
                )
            elif isinstance(callbacks, str):
                already_configured = callback_path in callbacks or "ceil_dlp_callback" in callbacks

            if already_configured:
                console.print(
                    f"callback already configured in [cyan]{litellm_config.name}[/cyan], skipping..."
                )
            else:
                # add callback
                if "litellm_settings" not in config:
                    config["litellm_settings"] = {}
                if "callbacks" not in config["litellm_settings"]:
                    config["litellm_settings"]["callbacks"] = []

                current_callbacks = config["litellm_settings"]["callbacks"]
                if isinstance(current_callbacks, str):
                    config["litellm_settings"]["callbacks"] = [current_callbacks]

                config["litellm_settings"]["callbacks"].append(callback_path)

                # write back
                with litellm_config.open("w") as f:
                    yaml.dump(config, f, default_flow_style=False, sort_keys=False)

                console.print(
                    f"updated [cyan]{litellm_config.name}[/cyan] to include ceil-dlp callback"
                )
                console.print(f"   [dim]Callback path:[/dim] [cyan]{callback_path}[/cyan]")
        except Exception as e:
            console.print(f"could not update [cyan]{litellm_config.name}[/cyan]: [red]{e}[/red]")
            console.print(
                f"   [dim]please manually add to litellm_settings.callbacks:[/dim] [cyan]{callback_path}[/cyan]"
            )

    console.print()
    console.print(
        Panel.fit("[bold green]installation complete![/bold green]", border_style="green")
    )

    next_steps = Text()
    next_steps.append("1. ", style="bold")
    next_steps.append("review and customize ", style="")
    next_steps.append(f"{ceil_dlp_config.name}", style="cyan")
    next_steps.append(" to suit your needs\n", style="")
    next_steps.append("2. ", style="bold")
    next_steps.append("run ", style="")
    next_steps.append(f"litellm --config {litellm_config} --port 4000", style="cyan")
    next_steps.append(" to start LiteLLM", style="")

    console.print("\n[bold]next steps:[/bold]\n")
    console.print(next_steps)
    console.print()


@app.command()
def remove(
    litellm_config: Path = typer.Argument(
        ...,
        help="Path to LiteLLM config.yaml file",
        exists=True,
        file_okay=True,
        dir_okay=False,
    ),
    remove_callback_file: bool = typer.Option(
        True,
        "--remove-callback-file/--keep-callback-file",
        help="Remove the ceil_dlp_callback.py wrapper file",
    ),
    remove_config_file: bool = typer.Option(
        False,
        "--remove-config-file/--keep-config-file",
        help="Remove the ceil-dlp.yaml config file",
    ),
    update_config: bool = typer.Option(
        True,
        "--update-config/--no-update-config",
        help="Automatically update LiteLLM config.yaml to remove ceil-dlp callback",
    ),
) -> None:
    """Remove ceil-dlp from the LiteLLM configuration.

    This command:
    1. Removes the callback from LiteLLM config.yaml
    2. Optionally removes the ceil_dlp_callback.py wrapper file
    3. Optionally removes the ceil-dlp.yaml config file
    """
    config_dir = litellm_config.parent
    callback_file = config_dir / "ceil_dlp_callback.py"
    ceil_dlp_config = config_dir / "ceil-dlp.yaml"

    console.print("\n[bold red]removing ceil-dlp[/bold red]")
    console.print(f"config directory: [cyan]{config_dir}[/cyan]")
    console.print(f"litellm config: [cyan]{litellm_config}[/cyan]\n")

    # remove callback from LiteLLM config if requested
    if update_config:
        try:
            with litellm_config.open() as f:
                config = yaml.safe_load(f) or {}

            callback_path = "ceil_dlp_callback.proxy_handler_instance"
            callbacks = config.get("litellm_settings", {}).get("callbacks", [])

            if not callbacks:
                console.print(
                    f"no callbacks configured in [cyan]{litellm_config.name}[/cyan], skipping..."
                )
            else:
                # normalize to list format
                if isinstance(callbacks, str):
                    callbacks = [callbacks]

                # find and remove ceil-dlp callbacks
                original_count = len(callbacks)
                callbacks = [
                    cb
                    for cb in callbacks
                    if callback_path not in str(cb) and "ceil_dlp_callback" not in str(cb)
                ]

                removed_count = original_count - len(callbacks)

                if removed_count > 0:
                    # update config
                    if len(callbacks) == 0:
                        # remove callbacks key if empty
                        if (
                            "litellm_settings" in config
                            and "callbacks" in config["litellm_settings"]
                        ):
                            del config["litellm_settings"]["callbacks"]
                            # remove litellm_settings if empty
                            if not config["litellm_settings"]:
                                del config["litellm_settings"]
                    else:
                        # convert back to string if only one callback remains
                        if len(callbacks) == 1:
                            config["litellm_settings"]["callbacks"] = callbacks[0]
                        else:
                            config["litellm_settings"]["callbacks"] = callbacks

                    # write back
                    with litellm_config.open("w") as f:
                        yaml.dump(config, f, default_flow_style=False, sort_keys=False)

                    console.print(
                        f"removed ceil-dlp callback from [cyan]{litellm_config.name}[/cyan]"
                    )
                else:
                    console.print(
                        f"ceil-dlp callback not found in [cyan]{litellm_config.name}[/cyan], skipping..."
                    )
        except Exception as e:
            console.print(f"could not update [cyan]{litellm_config.name}[/cyan]: [red]{e}[/red]")
            console.print(
                "   [dim]please manually remove ceil-dlp from litellm_settings.callbacks[/dim]"
            )

    # remove callback wrapper file if requested
    if remove_callback_file:
        if callback_file.exists():
            try:
                callback_file.unlink()
                console.print(f"removed [cyan]{callback_file.name}[/cyan]")
            except Exception as e:
                console.print(f"could not remove [cyan]{callback_file.name}[/cyan]: [red]{e}[/red]")
        else:
            console.print(f"[dim]{callback_file.name}[/dim] not found, skipping...")
    else:
        if callback_file.exists():
            console.print(f"[dim]keeping {callback_file.name}[/dim]")

    # remove ceil-dlp.yaml config file if requested
    if remove_config_file:
        if ceil_dlp_config.exists():
            try:
                ceil_dlp_config.unlink()
                console.print(f"removed [cyan]{ceil_dlp_config.name}[/cyan]")
            except Exception as e:
                console.print(
                    f"could not remove [cyan]{ceil_dlp_config.name}[/cyan]: [red]{e}[/red]"
                )
        else:
            console.print(f"[dim]{ceil_dlp_config.name}[/dim] not found, skipping...")
    else:
        if ceil_dlp_config.exists():
            console.print(f"[dim]keeping {ceil_dlp_config.name}[/dim]")

    console.print()
    console.print(Panel.fit("[bold green]removal complete![/bold green]", border_style="green"))


def main() -> None:
    """Main entry point for the CLI."""
    app()


if __name__ == "__main__":
    main()
