main.py

"""
main.py — Sentinel AI CLI entry point.

Authentication flow:
  1. On startup, check for a stored JWT token (~/.terminal_ai/token.json).
  2. If a valid (non-expired) token exists → greet user and go straight to chat.
  3. If no valid token → show the auth menu (Login / Register / Exit).
  4. Login / Register authenticate directly against the local database via
     auth_service.py — no HTTP server required.
  5. After successful authentication, the chat loop runs unchanged.

Chat workflow is delegated to chat_service.process_chat() so that
the same business logic is shared with the FastAPI layer.
"""

from prompts import SYSTEM_PROMPT  # noqa: F401 (imported for chat_service context)
from logger import logger
from memory import CURRENT_SESSION_ID
from utils import get_user_input, generate_session_id
from export import export_to_json, export_to_txt, export_to_markdown
from database.database import init_db, get_session
from database import repository
from database.utils import backup_database, integrity_check, get_db_size, health_check
from chat_service import process_chat
from token_store import save_token, delete_token, get_valid_token
import getpass
import sys
import os

# ---------------------------------------------------------------------------
# Database initialisation (always runs first, no auth required)
# ---------------------------------------------------------------------------
SESSION_ID = generate_session_id()
init_db()


# ---------------------------------------------------------------------------
# CLI auth banner helpers
# ---------------------------------------------------------------------------

def _print_auth_banner() -> None:
    """Print the startup menu."""
    print()
    print("=" * 44)
    print("   Sentinel AI  v5.0.0")
    print("=" * 44)
    print("  1.  Login")
    print("  2.  Register")
    print("  3.  Exit")
    print("=" * 44)


def _print_welcome(username: str) -> None:
    """Print the post-login greeting before the chat loop."""
    print()
    print("=" * 44)
    print(f"  Welcome, {username}")
    print("  Session Authenticated")
    print("=" * 44)
    print("  Type 'help' for available commands.")
    print("  Type 'exit' to quit.")
    print("=" * 44)
    print()


# ---------------------------------------------------------------------------
# Direct (local) auth — no HTTP, uses DB + auth_service directly
# ---------------------------------------------------------------------------

def _cli_login() -> str | None:
    """
    Prompt for credentials, authenticate against the local DB.
    Returns the authenticated username on success, None on failure.
    """
    from sqlmodel import select
    from database.models import User
    from auth_service import verify_password, create_access_token

    print()
    username = input("  Username : ").strip()
    if not username:
        print("  Username cannot be empty.")
        return None

    password = getpass.getpass("  Password : ")

    with get_session() as db:
        user = db.exec(select(User).where(User.username == username)).first()

    if not user or not verify_password(password, user.password_hash):
        print("  Invalid username or password.")
        return None

    token = create_access_token({"sub": user.username})
    save_token(token, user.username)
    return user.username


def _cli_register() -> bool:
    """
    Prompt for new credentials and register the user in the local DB.
    Returns True on success so the caller can redirect to login.
    """
    from sqlmodel import select
    from database.models import User
    from auth_service import hash_password

    print()
    username = input("  Choose a username : ").strip()
    if not username or len(username) < 3:
        print("  Username must be at least 3 characters.")
        return False

    password = getpass.getpass("  Choose a password : ")
    if not password or len(password) < 6:
        print("  Password must be at least 6 characters.")
        return False

    confirm = getpass.getpass("  Confirm password  : ")
    if password != confirm:
        print("  Passwords do not match.")
        return False

    with get_session() as db:
        existing = db.exec(select(User).where(User.username == username)).first()
        if existing:
            print(f"  Username '{username}' is already taken.")
            return False

        new_user = User(username=username, password_hash=hash_password(password))
        db.add(new_user)
        db.commit()

    print(f"  Account created. Please login as '{username}'.")
    return True


# ---------------------------------------------------------------------------
# Auth gate — runs before the chat loop
# ---------------------------------------------------------------------------

def run_auth_gate() -> str:
    """
    Ensure the user is authenticated before entering the chat loop.
    Returns the authenticated username.
    """
    # Fast path — reuse a stored, valid token
    stored = get_valid_token()
    if stored:
        username = stored["username"]
        _print_welcome(username)
        logger.info("CLI | Token reused for user: %s", username)
        return username

    # Slow path — show the menu
    while True:
        _print_auth_banner()
        choice = input("  Enter choice [1-3] : ").strip()

        if choice == "3":
            print("\n  Goodbye!\n")
            sys.exit(0)

        elif choice == "2":
            _cli_register()
            # After register, loop back to menu so user can login
            continue

        elif choice == "1":
            username = _cli_login()
            if username:
                logger.info("CLI | Login successful for user: %s", username)
                _print_welcome(username)
                return username
            else:
                print()  # blank line before re-showing menu

        else:
            print("  Please enter 1, 2, or 3.")


# ---------------------------------------------------------------------------
# Command handlers (unchanged from previous phases)
# ---------------------------------------------------------------------------

def show_help() -> None:
    print("\n--- Commands ---")
    print("  help    - Show available commands")
    print("  whoami  - Show currently logged in user")
    print("  logout  - Log out and return to menu")
    print("  clear   - Clear current session memory")
    print("  list    - List all saved sessions")
    print("  history - Show conversation history for this session")
    print("  stats   - Show database statistics")
    print("  export  - Export conversation to JSON, TXT, and MD")
    print("  backup  - Back up the database")
    print("  check   - Run database integrity check")
    print("  exit    - Quit the application")
    print("----------------")


def list_sessions() -> None:
    with get_session() as db:
        sessions = repository.list_sessions(db)

    if not sessions:
        print("\n  No sessions found in the database.")
        return

    print(f"\n--- Saved Sessions ({len(sessions)}) ---")
    for s in sessions:
        title = s.title or "(no title)"
        accessed = (
            s.last_accessed_at.strftime("%Y-%m-%d %H:%M")
            if s.last_accessed_at else "unknown"
        )
        marker = " <-- current" if s.session_id == CURRENT_SESSION_ID else ""
        print(f"  [{s.session_id}]  {title:<40}  last used: {accessed}{marker}")
    print("-------------------------------")


def show_history() -> None:
    with get_session() as db:
        messages = repository.get_messages(db, CURRENT_SESSION_ID)

    if not messages:
        print("\n  No messages in this session yet.")
        return

    print(f"\n--- Session History ({CURRENT_SESSION_ID}) ---")
    for msg in messages:
        role = msg.role.upper()
        display = msg.content if len(msg.content) <= 200 else msg.content[:197] + "..."
        print(f"\n  [{role}]\n  {display}")
    print("\n-------------------------------")


def show_stats() -> None:
    with get_session() as db:
        stats = repository.get_db_stats(db)

    db_size = get_db_size()
    hc = health_check()

    print("\n--- Database Statistics ---")
    print(f"  Sessions        : {stats['total_sessions']}")
    print(f"  Messages        : {stats['total_messages']}")
    print(f"  Incidents       : {stats['total_incidents']}")
    print(f"  Oldest session  : {stats['oldest_session']}")
    print(f"  Latest session  : {stats['latest_session']}")
    print(f"  Database size   : {db_size['human']}")
    print(f"  DB status       : {hc['status'].upper()}")
    print("-------------------------------")


def run_backup() -> None:
    try:
        path = backup_database()
        print(f"\n  Backup created: {path}")
    except Exception as error:
        print(f"\n  Backup failed: {error}")
        logger.exception("Backup failed: %s", error)


def run_integrity_check() -> None:
    result = integrity_check()
    if result["status"] == "ok":
        print("\n  Database integrity: OK")
    else:
        print("\n  Database integrity: ISSUES FOUND")
        for finding in result["findings"]:
            print(f"    - {finding}")


def clear_current_memory() -> None:
    from memory import clear_memory
    clear_memory()
    print("\nSession memory cleared.")


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------

logger.info("Application Started | Session : %s", CURRENT_SESSION_ID)

# Authenticate before doing anything else
CURRENT_USER = run_auth_gate()

# ---------------------------------------------------------------------------
# Chat loop (only reached after successful authentication)
# ---------------------------------------------------------------------------
try:
    while True:
        user_input = get_user_input()
        command = user_input.lower().strip()

        if command == "exit":
            logger.info("Application Closed")
            print("\nGoodbye!")
            break

        if command == "help":
            show_help()
            continue

        if command == "whoami":
            print(f"\n  Logged in as: {CURRENT_USER}")
            continue

        if command == "logout":
            delete_token()
            print("\n  Logged out. Restart the application to log in again.")
            logger.info("CLI | User logged out: %s", CURRENT_USER)
            break

        if command == "clear":
            clear_current_memory()
            continue

        if command == "list":
            list_sessions()
            continue

        if command == "history":
            show_history()
            continue

        if command == "stats":
            show_stats()
            continue

        if command == "backup":
            run_backup()
            continue

        if command == "check":
            run_integrity_check()
            continue

        if command == "export":
            user_choice = input(
                "\nExport formats (json, txt, md) [default: all]: "
            ).strip().lower()
            if user_choice == "json":
                export_to_json()
            elif user_choice == "txt":
                export_to_txt()
            elif user_choice == "md":
                export_to_markdown()
            else:
                export_to_json()
                export_to_txt()
                export_to_markdown()
            print("\nConversation exported successfully.")
            continue

        # --- Normal chat message: delegate entirely to chat_service ---
        with get_session() as db:
            result = process_chat(db, CURRENT_SESSION_ID, user_input)

        print("\nAssistant :", result.reply)

except KeyboardInterrupt:
    logger.info("Application Interrupted")
    print("\n\nGoodbye!")