#!/usr/bin/env python3
"""
norm-claude: Wrapper CLI for Claude Code with Norm Mode prompts
Designed for lawyers and coding novices who need situational awareness.

Usage:
    norm-claude                     Launch Claude Code in Norm Mode
    norm-claude doctor              Run health checks for setup
    norm-claude status              Print STATUS report schema
    norm-claude opinion             Print ERROR_OPINION schema
    norm-claude help                Print quickstart guide
    norm-claude init [directory]    Initialize Norm Mode project scaffold
    norm-claude --version           Show version
"""

import argparse
import subprocess
import sys
import os
from pathlib import Path

# Import version from package
try:
    from norm_claude import __version__ as VERSION
except ImportError:
    VERSION = "0.1.7"  # Fallback version

DEFAULT_CLAUDE_CMD = "claude"  # Configurable via --claude-cmd or env var

# Paths to prompt snippets - prioritize package resources over env var
PROMPT_DIR = None

# Try to load from installed package data first
try:
    from importlib.resources import files
    PROMPT_DIR = Path(str(files('norm_claude').joinpath('data/prompts')))
    # Verify it exists and has the required files
    if not (PROMPT_DIR / "NORM_MODE_SYSTEM_PROMPT.md").exists():
        PROMPT_DIR = None
except (ImportError, AttributeError, TypeError):
    try:
        # Python 3.7-3.8 fallback
        import pkg_resources
        PROMPT_DIR = Path(pkg_resources.resource_filename('norm_claude', 'data/prompts'))
        if not (PROMPT_DIR / "NORM_MODE_SYSTEM_PROMPT.md").exists():
            PROMPT_DIR = None
    except Exception:
        PROMPT_DIR = None

# Only use environment variable as fallback if package resources didn't work
if PROMPT_DIR is None:
    PROMPT_DIR = Path(os.getenv("NORM_MODE_PROMPTS",
                               Path.home() / "Documents" / "norm-mode" / "prompts"))

def print_banner():
    """Print Norm Mode banner when launching."""
    banner = """
    ╔══════════════════════════════════════════════════════════╗
    ║                                                          ║
    ║              NORM MODE for Claude Code                   ║
    ║                                                          ║
    ║  You're in a safe space. I'll always tell you:          ║
    ║    • What I'm about to do                               ║
    ║    • What changed and what did NOT change               ║
    ║    • How to undo anything                               ║
    ║                                                          ║
    ║  Type 'status' anytime to see where we are.             ║
    ║  Type 'opinion' after errors for a clear explanation.   ║
    ║                                                          ║
    ╚══════════════════════════════════════════════════════════╝
    """
    print(banner)


def launch_claude_code(claude_cmd, additional_args):
    """
    Launch Claude Code with Norm Mode system prompt.
    Reads NORM_MODE_SYSTEM_PROMPT.md and instructs user to load it.
    """
    prompt_file = PROMPT_DIR / "NORM_MODE_SYSTEM_PROMPT.md"

    if not prompt_file.exists():
        print(f"⚠️  Warning: System prompt not found at {prompt_file}")
        print(f"    Norm Mode will not be active. Run 'norm-claude init' to set up.")
        print()
        # Launch Claude anyway
        print(f"🚀 Launching Claude Code...")
        cmd = [claude_cmd] + additional_args
        try:
            subprocess.run(cmd)
        except FileNotFoundError:
            print(f"❌ Error: Claude Code command '{claude_cmd}' not found.")
            print()
            print("📚 Troubleshooting Steps:")
            print()
            print("1. Check if Claude Code is installed:")
            print("   ls -la ~/.local/bin/claude")
            print()
            print("2. If installed, reload your shell:")
            print("   source ~/.zshrc")
            print()
            print("3. If not installed, get Claude Code from:")
            print("   https://docs.anthropic.com/claude/docs/claude-code")
            print()
            print("4. Or set a custom command:")
            print(f"   norm-claude --claude-cmd /path/to/claude")
            print(f"   export CLAUDE_CMD=/path/to/claude")
            print()
            print("5. Run health check:")
            print("   norm-claude doctor")
            print()
            sys.exit(1)
        return

    # Launch Claude Code with NORM MODE system prompt
    print("📋 Loading NORM MODE system prompt...")
    print("✅ Launching Claude Code with NORM MODE activated")
    print()
    print("=" * 60)
    print("🎉 NORM MODE will be active when Claude opens")
    print("   No pasting required - it's automatic!")
    print("=" * 60)
    print()

    # Use --system-prompt-file flag to load prompt automatically
    cmd = [claude_cmd, '--system-prompt-file', str(prompt_file)] + additional_args

    try:
        subprocess.run(cmd)
    except FileNotFoundError:
        print(f"❌ Error: Claude Code command '{claude_cmd}' not found.")
        print()
        print("📚 Troubleshooting Steps:")
        print()
        print("1. Check if Claude Code is installed:")
        print("   ls -la ~/.local/bin/claude")
        print()
        print("2. If installed, reload your shell:")
        print("   source ~/.zshrc")
        print()
        print("3. If not installed, get Claude Code from:")
        print("   https://docs.anthropic.com/claude/docs/claude-code")
        print()
        print("4. Or set a custom command:")
        print(f"   norm-claude --claude-cmd /path/to/claude")
        print(f"   export CLAUDE_CMD=/path/to/claude")
        print()
        print("5. Run health check:")
        print("   norm-claude doctor")
        print()
        sys.exit(1)


def launch_brainstorm_mode(claude_cmd):
    """Launch Claude Code in brainstorming mode."""
    brainstorm_file = PROMPT_DIR / "BRAINSTORM_SNIPPET.md"

    if not brainstorm_file.exists():
        print(f"❌ BRAINSTORM_SNIPPET.md not found at: {brainstorm_file}")
        print("   Run 'norm-claude doctor' to check your setup.")
        sys.exit(1)

    try:
        brainstorm_prompt = brainstorm_file.read_text()
    except Exception as e:
        print(f"❌ Error reading brainstorm prompt: {e}")
        sys.exit(1)

    # Print brainstorm mode banner
    print("\n" + "="*60)
    print("🧠 BRAINSTORM MODE ACTIVATED")
    print("="*60)
    print("\nI'll help you develop a comprehensive specification for")
    print("your idea by asking one question at a time.")
    print("\nExit anytime by saying: 'exit brainstorm'")
    print("\n" + "="*60 + "\n")

    # Launch Claude with brainstorm prompt using --system-prompt-file
    print("Launching Claude Code in brainstorm mode...")
    print("\nProvide your idea and I'll start asking questions.\n")

    cmd = [claude_cmd, '--system-prompt-file', str(brainstorm_file)]

    try:
        subprocess.run(cmd)
    except FileNotFoundError:
        print(f"❌ Error: Claude Code command '{claude_cmd}' not found.")
        print()
        print("📚 Troubleshooting Steps:")
        print()
        print("1. Check if Claude Code is installed:")
        print("   ls -la ~/.local/bin/claude")
        print()
        print("2. If installed, reload your shell:")
        print("   source ~/.zshrc")
        print()
        print("3. Run health check:")
        print("   norm-claude doctor")
        print()
        sys.exit(1)


def print_status_snippet():
    """Print the STATUS schema snippet for user reference."""
    snippet_file = PROMPT_DIR / "STATUS_SNIPPET.md"
    if snippet_file.exists():
        try:
            print(snippet_file.read_text())
        except Exception as e:
            print(f"❌ Error reading STATUS snippet: {e}")
            sys.exit(1)
    else:
        print(f"❌ STATUS snippet not found at: {snippet_file}")
        print("   Run 'norm-claude init' to set up Norm Mode.")
        sys.exit(1)


def print_opinion_snippet():
    """Print the ERROR_OPINION schema snippet for user reference."""
    snippet_file = PROMPT_DIR / "ERROR_OPINION_SNIPPET.md"
    if snippet_file.exists():
        try:
            print(snippet_file.read_text())
        except Exception as e:
            print(f"❌ Error reading ERROR_OPINION snippet: {e}")
            sys.exit(1)
    else:
        print(f"❌ ERROR_OPINION snippet not found at: {snippet_file}")
        print("   Run 'norm-claude init' to set up Norm Mode.")
        sys.exit(1)


def print_status_lite_snippet():
    """Print the status-lite schema snippet for user reference."""
    snippet_file = PROMPT_DIR / "STATUS_LITE_SNIPPET.md"
    if snippet_file.exists():
        try:
            print(snippet_file.read_text())
        except Exception as e:
            print(f"❌ Error reading status-lite snippet: {e}")
            sys.exit(1)
    else:
        print(f"❌ Status-lite snippet not found at: {snippet_file}")
        print("   Run 'norm-claude init' to set up Norm Mode.")
        sys.exit(1)


def print_help_guide():
    """Print the quickstart guide."""
    guide_file = PROMPT_DIR / "QUICKSTART_FOR_LAWYERS.md"
    if guide_file.exists():
        try:
            print(guide_file.read_text())
        except Exception as e:
            print(f"❌ Error reading quickstart guide: {e}")
            sys.exit(1)
    else:
        print("""
NORM MODE QUICKSTART

Commands:
  norm-claude                Launch Claude Code in Norm Mode
  norm-claude status         Show STATUS report schema
  norm-claude opinion        Show ERROR_OPINION schema
  norm-claude help           Show this guide
  norm-claude init [dir]     Initialize Norm Mode project

In Claude, you can type:
  status      Get current situational awareness
  opinion     Get error analysis
  risk        Understand risk level of last action
  explain     Define technical terms

Run 'norm-claude init' to set up full documentation.
        """)


def init_project(directory):
    """
    Initialize Norm Mode project scaffold.
    Creates directory structure and copies templates.
    """
    target = Path(directory).expanduser().resolve()

    # Create subdirectories
    try:
        target.mkdir(parents=True, exist_ok=True)
        (target / "prompts").mkdir(exist_ok=True)
        (target / "examples").mkdir(exist_ok=True)
        (target / "cli").mkdir(exist_ok=True)
    except Exception as e:
        print(f"❌ Error creating directories: {e}")
        sys.exit(1)

    print(f"✅ Created Norm Mode scaffold at: {target}")
    print()

    # Check if running from installed location vs development
    script_dir = Path(__file__).parent
    source_prompts = script_dir.parent / "prompts"
    source_examples = script_dir.parent / "examples"

    # Copy prompt files if they exist
    if source_prompts.exists():
        print("📋 Copying prompt files...")
        try:
            import shutil
            for prompt_file in source_prompts.glob("*.md"):
                dest = target / "prompts" / prompt_file.name
                shutil.copy2(prompt_file, dest)
                print(f"   ✓ {prompt_file.name}")
        except Exception as e:
            print(f"   ⚠️  Could not copy prompts: {e}")

    # Copy example files if they exist
    if source_examples.exists():
        print("📚 Copying example scenarios...")
        try:
            import shutil
            for example_file in source_examples.glob("*.md"):
                dest = target / "examples" / example_file.name
                shutil.copy2(example_file, dest)
                print(f"   ✓ {example_file.name}")
        except Exception as e:
            print(f"   ⚠️  Could not copy examples: {e}")

    print()
    print("Next steps:")
    print(f"  1. Set environment variable:")
    print(f"     export NORM_MODE_PROMPTS=\"{target / 'prompts'}\"")
    print(f"     (Add this to ~/.zshrc or ~/.bash_profile to make permanent)")
    print()
    print(f"  2. Launch Norm Mode:")
    print(f"     norm-claude")
    print()
    print(f"  3. View documentation:")
    print(f"     norm-claude help")


def run_doctor(claude_cmd):
    """
    Run health checks to verify Norm Mode setup.
    Checks Claude installation, PATH, environment variables, and prompt files.
    """
    print("🏥 Running Norm Mode Health Check...")
    print()

    issues = []
    warnings = []

    # Check 1: Claude Code executable exists
    print("1. Checking Claude Code installation...")
    common_paths = [
        Path.home() / ".local" / "bin" / "claude",
        Path("/usr/local/bin/claude"),
        Path("/opt/homebrew/bin/claude"),
    ]

    claude_found = False
    claude_location = None

    # Check if the configured claude_cmd exists
    try:
        result = subprocess.run([claude_cmd, "--version"],
                              capture_output=True,
                              text=True,
                              timeout=5)
        if result.returncode == 0:
            claude_found = True
            claude_location = claude_cmd
            print(f"   ✅ Claude Code found: {claude_cmd}")
            print(f"      Version: {result.stdout.strip()}")
        else:
            issues.append(f"Claude command '{claude_cmd}' exists but --version failed")
    except FileNotFoundError:
        issues.append(f"Claude command '{claude_cmd}' not found in PATH")
        print(f"   ❌ Claude command '{claude_cmd}' not found")

        # Check common locations
        print("   🔍 Checking common installation locations...")
        for path in common_paths:
            if path.exists():
                print(f"      Found: {path}")
                claude_location = str(path)
                warnings.append(f"Claude found at {path} but not in PATH")
    except subprocess.TimeoutExpired:
        warnings.append(f"Claude command '{claude_cmd}' timed out")

    print()

    # Check 2: NORM_MODE_PROMPTS environment variable
    print("2. Checking NORM_MODE_PROMPTS environment variable...")
    norm_prompts_env = os.getenv("NORM_MODE_PROMPTS")
    if norm_prompts_env:
        if not Path(norm_prompts_env).exists():
            issues.append(f"NORM_MODE_PROMPTS points to non-existent directory: {norm_prompts_env}")
            print(f"   ❌ NORM_MODE_PROMPTS is set but directory does not exist: {norm_prompts_env}")
        else:
            print(f"   ✅ NORM_MODE_PROMPTS is set: {norm_prompts_env}")
    else:
        # Check if package resources are being used (which is the normal case for PyPI installs)
        if "site-packages" in str(PROMPT_DIR):
            print(f"   ✅ Using package resources (no environment variable needed)")
            print(f"      Prompts loaded from: {PROMPT_DIR}")
        else:
            # Only warn if package resources aren't working and env var is missing
            warnings.append("NORM_MODE_PROMPTS not set and package resources not found")
            print(f"   ⚠️  NORM_MODE_PROMPTS not set")
            print(f"      Using: {PROMPT_DIR}")

    print()

    # Check 3: System prompt file exists
    print("3. Checking system prompt file...")
    prompt_file = PROMPT_DIR / "NORM_MODE_SYSTEM_PROMPT.md"
    if prompt_file.exists():
        print(f"   ✅ System prompt found: {prompt_file}")
        file_size = prompt_file.stat().st_size
        print(f"      Size: {file_size:,} bytes")
    else:
        issues.append(f"System prompt not found at {prompt_file}")
        print(f"   ❌ System prompt not found: {prompt_file}")

    print()

    # Check 4: Other required prompt files
    print("4. Checking other prompt files...")
    required_files = ["STATUS_SNIPPET.md", "STATUS_LITE_SNIPPET.md", "ERROR_OPINION_SNIPPET.md", "BRAINSTORM_SNIPPET.md", "QUICKSTART_FOR_LAWYERS.md"]
    for filename in required_files:
        file_path = PROMPT_DIR / filename
        if file_path.exists():
            print(f"   ✅ {filename}")
        else:
            warnings.append(f"Optional file missing: {filename}")
            print(f"   ⚠️  {filename} (optional)")

    print()

    # Check 5: Shell configuration
    print("5. Checking shell configuration...")
    shell_config = Path.home() / ".zshrc"
    if shell_config.exists():
        content = shell_config.read_text()

        # Check PATH configuration
        if ".local/bin" in content or "~/.local/bin" in content:
            print(f"   ✅ ~/.local/bin in PATH")
        else:
            warnings.append("~/.local/bin might not be in PATH")
            print(f"   ⚠️  ~/.local/bin not found in {shell_config}")
            print(f"      Add: export PATH=\"$HOME/.local/bin:$PATH\"")
    else:
        warnings.append(f"{shell_config} not found")
        print(f"   ⚠️  {shell_config} not found")

    print()
    print("=" * 60)

    # Summary
    if not issues and not warnings:
        print("✅ All checks passed! Norm Mode is ready to use.")
        print()
        print("Run 'norm-claude' to launch Claude Code in Norm Mode.")
    elif issues:
        print("❌ Critical Issues Found:")
        for issue in issues:
            print(f"   • {issue}")
        print()
        if warnings:
            print("⚠️  Warnings:")
            for warning in warnings:
                print(f"   • {warning}")
            print()
        print("Fix the critical issues above, then run 'norm-claude doctor' again.")
    else:
        print("⚠️  Warnings Found:")
        for warning in warnings:
            print(f"   • {warning}")
        print()
        print("Norm Mode should work, but you may want to address the warnings above.")
        print("Run 'norm-claude' to launch Claude Code in Norm Mode.")

    print()


def main():
    parser = argparse.ArgumentParser(
        description="Norm Mode wrapper for Claude Code",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  norm-claude                        Launch Claude Code in Norm Mode
  norm-claude doctor                 Run health checks
  norm-claude status                 Print full STATUS schema (120-160 words)
  norm-claude status-lite            Print status-lite schema (40-60 words)
  norm-claude opinion                Print ERROR_OPINION schema
  norm-claude brainstorm             Launch interactive brainstorming session
  norm-claude help                   Show quickstart guide
  norm-claude init ~/Documents/norm-mode    Initialize project

Environment Variables:
  NORM_MODE_PROMPTS    Path to prompts directory
  CLAUDE_CMD           Claude Code command (default: claude)

For more information, visit: https://github.com/yourorg/norm-mode
        """
    )

    parser.add_argument("--version", action="version", version=f"norm-claude {VERSION}")
    parser.add_argument("--claude-cmd",
                        default=os.getenv("CLAUDE_CMD", DEFAULT_CLAUDE_CMD),
                        help="Claude Code command to run (default: claude)")

    subparsers = parser.add_subparsers(dest="command", help="Commands")

    # Subcommands
    subparsers.add_parser("status", help="Print full STATUS report schema (120-160 words)")
    subparsers.add_parser("status-lite", help="Print status-lite report schema (40-60 words)")
    subparsers.add_parser("opinion", help="Print ERROR_OPINION schema")
    subparsers.add_parser("brainstorm", help="Launch Claude Code in brainstorming mode")
    subparsers.add_parser("help", help="Print quickstart guide")
    subparsers.add_parser("doctor", help="Run health checks for Norm Mode setup")

    init_parser = subparsers.add_parser("init", help="Initialize Norm Mode project")
    init_parser.add_argument(
        "directory",
        nargs="?",
        default=str(Path.home() / "Documents" / "norm-mode"),
        help="Directory to initialize (default: ~/Documents/norm-mode)"
    )

    # Parse known args to allow passing through to Claude
    args, unknown = parser.parse_known_args()

    # Route to appropriate handler
    if args.command == "status":
        print_status_snippet()
    elif args.command == "status-lite":
        print_status_lite_snippet()
    elif args.command == "opinion":
        print_opinion_snippet()
    elif args.command == "brainstorm":
        launch_brainstorm_mode(args.claude_cmd)
    elif args.command == "help":
        print_help_guide()
    elif args.command == "doctor":
        run_doctor(args.claude_cmd)
    elif args.command == "init":
        init_project(args.directory)
    else:
        # Default: launch Claude with Norm Mode
        print_banner()
        launch_claude_code(args.claude_cmd, unknown)


if __name__ == "__main__":
    main()
