#!/usr/bin/env python3
"""meshulash-guard Demo Chatbot

Reference implementation showing how to integrate meshulash-guard
into a Gemini-powered LLM application.

Usage:
    export GEMINI_API_KEY="..."
    export MESHULASH_API_KEY="..."
    export MESHULASH_TENANT_ID="..."
    export MESHULASH_SERVER_URL="https://app.meshulash.ai"
    python demo/chatbot.py
"""

import os
import sys

from google import genai
from google.genai import types

from meshulash_guard import (
    Guard, Action, Condition,
    PIIScanner, PIILabel,
    ToxicityScanner, ToxicityLabel,
    TopicScanner, TopicLabel,
    CyberScanner, CyberLabel,
    JailbreakScanner, JailbreakLabel,
    ScanResult,
    AuthError, ConnectionError, ServerError,
)


# ---------------------------------------------------------------------------
# Environment
# ---------------------------------------------------------------------------

def _require_env(name: str) -> str:
    """Read an environment variable, exit with a clear message if missing."""
    val = os.environ.get(name)
    if not val:
        print(f"Error: {name} environment variable is not set.", file=sys.stderr)
        sys.exit(1)
    return val


GEMINI_API_KEY = _require_env("GEMINI_API_KEY")
MESHULASH_API_KEY = _require_env("MESHULASH_API_KEY")
MESHULASH_TENANT_ID = _require_env("MESHULASH_TENANT_ID")
MESHULASH_SERVER_URL = _require_env("MESHULASH_SERVER_URL")


# ---------------------------------------------------------------------------
# Presets
# ---------------------------------------------------------------------------

PRESETS = {
    "1": {
        "name": "PII Protection",
        "description": "Replaces PII (names, emails, phones, addresses) with placeholders",
        "scanners": [
            PIIScanner(labels=[PIILabel.PII], action=Action.REPLACE),
        ],
    },
    "2": {
        "name": "PII + Overrides",
        "description": "Replace most PII, but LOG secrets rather than replace them",
        "scanners": [
            PIIScanner(
                labels=[PIILabel.PII, PIILabel.SECRETS],
                action=Action.REPLACE,
                overrides={PIILabel.SECRETS: Action.LOG},
            ),
        ],
    },
    "3": {
        "name": "Content Moderation",
        "description": "Blocks toxic content and restricts sensitive topics",
        "scanners": [
            ToxicityScanner(
                labels=[ToxicityLabel.TOXICITY, ToxicityLabel.HATE_SPEECH],
                action=Action.BLOCK,
            ),
            TopicScanner(
                labels=[TopicLabel.SEXUAL_EXPLICIT, TopicLabel.IDENTITY_ATTACK],
                action=Action.BLOCK,
            ),
        ],
    },
    "4": {
        "name": "Cyber Threat Detection",
        "description": "Blocks prompts containing cyber attack content",
        "scanners": [
            CyberScanner(labels=[CyberLabel.ALL], action=Action.BLOCK),
        ],
    },
    "5": {
        "name": "Full Guard",
        "description": "All scanners active: PII replaced, toxic/cyber blocked",
        "scanners": [
            PIIScanner(labels=[PIILabel.ALL], action=Action.REPLACE),
            ToxicityScanner(labels=[ToxicityLabel.ALL], action=Action.BLOCK),
            CyberScanner(labels=[CyberLabel.ALL], action=Action.BLOCK),
            # Forward-compatible: activates when server deploys jailbreak TC head
            JailbreakScanner(labels=[JailbreakLabel.ALL], action=Action.BLOCK),
        ],
    },
}


# ---------------------------------------------------------------------------
# Output helpers
# ---------------------------------------------------------------------------

def _print_scan_result(direction: str, result: ScanResult) -> None:
    """Print scan result in a developer-realistic format.

    Shows status line always, plus triggered scanner details and placeholder
    counts when relevant. Mirrors what a real developer would log in their app.
    """
    status_tag = f"[{result.status.upper()}]"
    print(f"\n-- {direction} SCAN {status_tag} --")

    if result.risk_categories:
        print(f"   Risk categories: {', '.join(result.risk_categories)}")

    for scanner_id, sr in result.scanners.items():
        if sr.triggered:
            detections_summary = ", ".join(
                f"{d.label}({d.action})" for d in sr.detections
            )
            print(f"   {scanner_id}: triggered | score={sr.score:.2f} | {detections_summary}")

    if result.placeholders:
        print(f"   Placeholders created: {len(result.placeholders)}")


# ---------------------------------------------------------------------------
# Chat turn
# ---------------------------------------------------------------------------

def chat_turn(guard: Guard, chat_session, scanners: list, user_input: str) -> None:
    """Execute one turn of the chatbot loop with full scan/gate/deanonymize workflow.

    Workflow:
        1. scan_input()  - scan the user message
        2. Gate on BLOCK - if blocked, do not forward to Gemini
        3. send_message() - send processed (possibly redacted) text to Gemini
        4. scan_output() - scan the LLM response
        5. Gate on BLOCK - if blocked, do not return to user
        6. deanonymize() - restore original PII values in final text
    """
    # 1. Scan input
    input_result = guard.scan_input(user_input, scanners=scanners)
    _print_scan_result("INPUT", input_result)

    # 2. Gate on block — never forward blocked input to the LLM
    if input_result.status == "blocked":
        print("[BLOCKED] Input was blocked. Not forwarding to LLM.")
        return

    # 3. Send processed (possibly redacted) text to Gemini
    #    Always use processed_text — not the raw user_input — to ensure
    #    replaced PII is sent as placeholders, not original values.
    text_to_send = input_result.processed_text
    gemini_response = chat_session.send_message(text_to_send)
    raw_llm_text = gemini_response.text

    # 4. Scan LLM output
    output_result = guard.scan_output(raw_llm_text, scanners=scanners)
    _print_scan_result("OUTPUT", output_result)

    # 5. Gate on block — do not return blocked LLM output to user
    if output_result.status == "blocked":
        print("[BLOCKED] LLM response blocked by output scanner.")
        return

    # 6. Deanonymize: replace placeholders with original values
    #    Chain: gemini_response.text -> scan_output() -> processed_text -> deanonymize() -> final
    #    No-op if no PII was replaced (vault is empty or no placeholders in text).
    final_text = guard.deanonymize(output_result.processed_text)

    print(f"\nAssistant: {final_text}\n")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    """Run the interactive chatbot demo."""
    print("=" * 60)
    print("  meshulash-guard Demo Chatbot")
    print("=" * 60)
    print()
    print("Select a security preset:")
    print()
    for key, preset in PRESETS.items():
        print(f"  {key}. {preset['name']}")
        print(f"     {preset['description']}")
        print()

    choice = input("Preset number: ").strip()
    if choice not in PRESETS:
        print(f"Invalid choice '{choice}'. Must be one of: {', '.join(PRESETS.keys())}.")
        sys.exit(1)

    preset = PRESETS[choice]
    print(f"\nActive preset: {preset['name']}")
    print(f"Description:   {preset['description']}")
    print()

    # Initialize Guard (validates credentials at startup)
    print("Initializing meshulash-guard...")
    guard = Guard(
        api_key=MESHULASH_API_KEY,
        tenant_id=MESHULASH_TENANT_ID,
        server_url=MESHULASH_SERVER_URL,
    )
    print("Guard initialized.\n")

    # Initialize Gemini chat session with persistent conversation history
    gemini_client = genai.Client(api_key=GEMINI_API_KEY)
    chat_session = gemini_client.chats.create(
        model="gemini-2.5-flash",
        config=types.GenerateContentConfig(
            system_instruction="You are a helpful assistant.",
        ),
    )

    scanners = preset["scanners"]

    print("Chat started. Type your message. Press Ctrl+C or Ctrl+D to exit.")
    print("-" * 60)
    print()

    while True:
        try:
            user_input = input("You: ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\nGoodbye.")
            break

        if not user_input:
            continue

        try:
            chat_turn(guard, chat_session, scanners, user_input)
        except (AuthError, ConnectionError, ServerError) as e:
            print(f"[ERROR] SDK error: {e}")


if __name__ == "__main__":
    main()
