from __future__ import annotations

import argparse
import json
import logging
import os
import shutil
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable

if sys.version_info >= (3, 11):
    import tomllib
else:
    import tomli as tomllib  # type: ignore


is_windows = sys.platform == "win32"
is_linux = sys.platform == "linux"
is_macos = sys.platform == "darwin"

logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger(__name__)
cwd = Path.cwd()

_BUILD_COMMANDS = ["uv", "poetry", "hatch"]


def parse_pyproject_toml(directory: Path) -> dict:
    """Parse pyproject.toml file in directory and return project data."""
    project_toml = directory / "pyproject.toml"
    if not project_toml.is_file():
        logger.error(f"No pyproject.toml found in {directory}")
        return {}

    with project_toml.open("rb") as f:
        return tomllib.load(f)


def _get_build_command_from_toml(directory: Path) -> str | None:
    """Get build command from pyproject.toml."""
    logger.debug(f"Parsing pyproject.toml in {directory}")

    project_data = parse_pyproject_toml(directory)
    if not project_data:
        return None

    if "build-system" in project_data:
        build_system = project_data["build-system"]
        if "build-backend" in build_system:
            build_backend = build_system["build-backend"]
            if build_backend.startswith("poetry."):
                return "poetry"
            elif build_backend.startswith("hatchling."):
                return "hatch"
            else:
                logger.error(f"Unknown build-backend: {build_backend}")
                return None

    logger.error("No `build-system` or `build-backend` found in pyproject.toml: ")
    logger.error(json.dumps(project_data, indent=2, ensure_ascii=False, sort_keys=True))
    return None


def _get_build_command(directory: Path):
    """Get build command from directory."""
    project_path = directory / "pyproject.toml"
    if project_path.is_file():
        logger.debug(f"Found pyproject.toml in {directory}")
        return _get_build_command_from_toml(directory)

    for command in _BUILD_COMMANDS:
        if shutil.which(command):
            logger.debug(f"Found build command: {command}")
            return command
    logger.error(f"No build command found in {directory}")
    sys.exit(1)
    return None


@dataclass
class Command:
    name: str
    alias: str
    cmds: list[str] | Callable[..., Any] | None = None


def _clean(root_dir: Path):
    _run_command(["rm", "-rf", "dist", "build", "*.egg-info"], root_dir)


BUILD_COMMAND = _get_build_command(cwd) or ""
_COMMANDS = [
    Command(name="build", alias="b", cmds=[BUILD_COMMAND, "build"]),
    Command(name="bumpversion", alias="bump", cmds=["bumpversion", "patch"]),
    Command(name="clean", alias="c", cmds=_clean),
    Command(name="publish", alias="p"),  # No preset commands
    Command(name="test", alias="t", cmds=lambda: os.system("pytest")),
    Command(name="test-benchmark", alias="tb", cmds=lambda: os.system("pytest -m benchmark")),
    Command(name="test-coverage", alias="tc", cmds=lambda: os.system("pytest --cov=sfi")),
    Command(name="token", alias="tk", cmds=lambda: _set_token(BUILD_COMMAND)),
]
_COMMAND_DICT = {command.name: command for command in _COMMANDS}
_COMMAND_DICT.update({command.alias: command for command in _COMMANDS})
_CHOICES = [command.alias for command in _COMMANDS]
_CHOICES.extend([command.name for command in _COMMANDS])


def main():
    parser = argparse.ArgumentParser(description="Make Python")
    parser.add_argument("command", type=str, choices=_CHOICES, help=f"Command to run, options: {_CHOICES}")
    parser.add_argument("--debug", "-d", action="store_true", help="Enable debug mode")

    args = parser.parse_args()
    if args.debug:
        logger.setLevel(logging.DEBUG)

    build_command = _get_build_command(cwd) or ""
    logger.info(f"Using build command: {build_command}")
    command = _COMMAND_DICT.get(args.command)
    if command:
        if callable(command.cmds):
            command.cmds()
        elif isinstance(command.cmds, list):
            _run_command(command.cmds, cwd)
        else:
            logger.debug("No preset commands found")
    else:
        logger.error(f"Unknown command: {args.command}")
        sys.exit(1)

    if args.command in {"publish", "p"}:
        if not _check_pypi_token(build_command):
            _set_token(build_command)
        _run_command([build_command, "publish"], cwd)


def _set_token(build_command: str, show_header: bool = True) -> None:
    """Set PyPI token for the specified build command."""
    if show_header:
        logger.info(f"Setting PyPI token for {build_command}...")

    if build_command.lower() not in _BUILD_COMMANDS:
        logger.error(f"Unknown build command: {build_command}")
        logger.error(f"Please use `{'/'.join(_BUILD_COMMANDS)}`")
        sys.exit(1)

    token = input("Enter your PyPI token (leave empty to cancel): ").strip()
    if not token:
        logger.info("Invalid token, cancelled.")
        return

    if build_command == "uv":
        _set_uv_token(token)
    elif build_command == "poetry":
        _set_poetry_token(token)
    elif build_command == "hatch":
        _set_hatch_token(token)
    else:
        logger.error(f"Unknown build command: {build_command}")
        sys.exit(1)


def _set_uv_token(token: str) -> None:
    """Set PyPI token for uv."""
    _write_to_env_file("UV_PUBLISH_TOKEN", token)

    # Write to `uv.toml`
    config_path = Path.home() / ".config" / "uv" / "uv.toml"
    config_path.parent.mkdir(parents=True, exist_ok=True)
    content = f"""[publish]
token = "{token}"
"""
    config_path.write_text(content)
    logger.info(f"Token saved to {config_path}")


def _set_poetry_token(token: str) -> None:
    """Set PyPI token for poetry."""
    _write_to_env_file("POETRY_PYPI_TOKEN_PYPI", token)
    _run_command(["poetry", "config", "pypi-token.pypi", token], cwd)
    logger.info("Token saved to Poetry configuration.")


def _set_hatch_token(token: str) -> None:
    """Set PyPI token for hatch."""
    pypirc_path = Path.home() / ".pypirc"
    pypirc_content = f"""[pypi]
repository = https://upload.pypi.org/legacy/
username = __token__
password = {token}
"""
    pypirc_path.write_text(pypirc_content)
    logger.info(f"Token saved to {pypirc_path}")


def _check_pypi_token(build_command: str) -> bool:
    """Check if PyPI token is configured before publishing."""
    logger.info("Checking PyPI token configuration...")

    token_found = False
    if build_command == "uv":
        # Check for uv publish token
        token_env_vars = ["UV_PUBLISH_TOKEN", "PYPI_API_TOKEN"]
        for var in token_env_vars:
            if os.getenv(var):
                logger.info(f"Found PyPI token in environment variable: {var}")
                token_found = True
                break

        # Check for config file
        config_path = Path.home() / ".config" / "uv" / "uv.toml"
        if config_path.exists():
            logger.info(f"Found uv config file: {config_path}")
            token_found = True

    elif build_command == "poetry":
        # Check for poetry token
        if os.getenv("POETRY_PYPI_TOKEN_PYPI"):
            logger.info("Found PyPI token in POETRY_PYPI_TOKEN_PYPI environment variable")
            token_found = True

        # Check for poetry config
        result = subprocess.run(
            ["poetry", "config", "pypi-token.pypi"],
            capture_output=True,
            text=True,
        )
        if result.stdout.strip() and result.stdout.strip() != "None":
            logger.info("Found PyPI token in Poetry configuration")
            token_found = True

    elif build_command == "hatch":
        # Check for .pypirc
        pypirc_path = Path.home() / ".pypirc"
        if pypirc_path.exists():
            logger.info(f"Found .pypirc file: {pypirc_path}")
            token_found = True

    return token_found


def _run_command(cmd: list[str], directory: Path) -> None:
    """Run a command in the specified directory."""
    logger.debug(f"Running command: {' '.join(cmd)}")
    try:
        result = subprocess.run(cmd, cwd=directory, capture_output=True, text=True)
    except subprocess.CalledProcessError as e:
        logger.error(f"Command failed with exit code {e.returncode}")
        sys.exit(e.returncode)

    if result.stdout:
        print(result.stdout)
    if result.stderr:
        print(result.stderr, file=sys.stderr)


def _write_to_env_file(key: str, value: str) -> None:
    """Write key-value pair to environment file."""
    if is_windows:
        subprocess.run(["setx", key, value], shell=True)
    else:
        _write_to_shell_config(f"export {key}='{value}'")


def _get_shell_config_path() -> Path:
    """Get the appropriate shell config file based on the current shell."""
    # Try to detect the shell
    shell = os.getenv("SHELL", "")
    if "zsh" in shell:
        return Path.home() / ".zshrc"
    else:
        # Default to .bashrc
        return Path.home() / ".bashrc"


def _write_to_shell_config(content: str) -> None:
    """Write content to ~/.bashrc, replacing existing entries if they exist.

    Deprecated: Use _write_to_shell_config instead.
    """
    config_path = _get_shell_config_path()
    if not config_path.exists():
        logger.warning(f"{config_path} does not exist, creating it...")
        config_path.touch()

    # Extract the variable name from the export statement
    # Expected format: export VARIABLE_NAME=value
    var_name = None
    for line in content.strip().split("\n"):
        if line.startswith("export ") and "=" in line:
            var_name = line.split("=")[0].replace("export ", "").strip()
            break

    if not var_name:
        logger.error("Invalid export statement format. Expected: export VARIABLE_NAME=value")
        return

    # Read existing content
    existing_lines = config_path.read_text(encoding="utf-8").split("\n")

    # Find and remove existing export statements for this variable
    new_lines = []
    found_existing = False
    for line in existing_lines:
        # Check if this line exports the same variable
        if line.strip().startswith(f"export {var_name}=") or line.strip().startswith(f"export {var_name} ="):
            found_existing = True
            continue
        new_lines.append(line)

    if found_existing:
        logger.info(f"Found existing export statement for {var_name}, replacing it...")

    # Add new content
    new_lines.append(content.strip())

    # Write back to file
    config_path.write_text("\n".join(new_lines), encoding="utf-8")

    logger.info(f"Content written to {config_path}")
    logger.info(f"Run `source {config_path}` to apply the changes")
    logger.info(f"Run `cat {config_path}` to view the content")
    logger.info(f"Run `cat {config_path} | grep 'export'` to view the exported variables")
