#!/usr/bin/env python3
"""
AI Dev Vault Health Check -- Validates your vault setup is complete and correct.

Checks:
  - Directory structure exists
  - All agent definitions have required sections
  - Hooks are configured and executable
  - Learning files have correct format
  - CLAUDE.md is present and under 300 lines
  - Python dependencies installed
  - vault-search index is fresh
  - Git is configured
  - No secrets in tracked files

Usage:
    health-check                    # run all checks
    health-check --fix              # auto-fix what's possible
    health-check --quiet            # only show failures
    health-check --json             # output as JSON
    health-check --version          # show version
"""

import sys
import os
import re
import json
import subprocess
from pathlib import Path

__version__ = "1.0.0"

# --- Configuration ---
VAULT_ROOT = Path(os.environ.get("VAULT_ROOT", ".")).resolve()
PROJECT_DIR = Path(os.environ.get("PROJECT_DIR", ".")).resolve()

# ANSI colors
class C:
    GREEN = "\033[92m"
    YELLOW = "\033[93m"
    RED = "\033[91m"
    BLUE = "\033[94m"
    BOLD = "\033[1m"
    DIM = "\033[2m"
    RESET = "\033[0m"

    @staticmethod
    def disable():
        C.GREEN = C.YELLOW = C.RED = C.BLUE = C.BOLD = C.DIM = C.RESET = ""


# ============================================================
# Check Functions
# ============================================================

def check_directory_structure():
    """Verify required directories exist."""
    issues = []
    required_dirs = [
        ".claude/agents",
        ".claude/skills",
        ".claude/hooks",
        "agent-learnings",
        "handoffs",
        "docs/status",
        "docs",
    ]

    for d in required_dirs:
        path = PROJECT_DIR / d
        if not path.exists():
            issues.append(f"Missing directory: {d}")

    optional_dirs = [
        ".session",
        "agent-reports",
        ".codebase-map",
    ]
    warnings = []
    for d in optional_dirs:
        path = PROJECT_DIR / d
        if not path.exists():
            warnings.append(f"Optional directory missing: {d} (created automatically when needed)")

    return {
        "name": "Directory Structure",
        "passed": len(issues) == 0,
        "issues": issues,
        "warnings": warnings,
    }


def check_agents():
    """Verify agent definitions have required sections."""
    issues = []
    warnings = []
    agents_dir = PROJECT_DIR / ".claude" / "agents"

    if not agents_dir.exists():
        return {"name": "Agent Definitions", "passed": False,
                "issues": ["No .claude/agents/ directory"], "warnings": []}

    # Agents may use different header names, so check for common variants
    required_sections_variants = [
        ["Process", "Your Process", "Workflow", "Steps", "How to", "Instructions", "Your Job", "Your Role",
         "Before You Start", "While You Work", "What You Do", "What You Check", "Audit Process", "Validation Process", "Scope"],
        ["Rules", "Your Rules", "Constraints", "Guidelines", "Behavioral Rules"],
    ]
    recommended_sections = ["Output Format", "System Cross-Reference"]

    agent_files = list(agents_dir.glob("*.md"))
    if not agent_files:
        issues.append("No agent definitions found in .claude/agents/")

    for agent_file in agent_files:
        content = agent_file.read_text(encoding="utf-8", errors="ignore")

        # Check frontmatter
        if not content.startswith("---"):
            issues.append(f"{agent_file.name}: Missing YAML frontmatter")

        # Check required sections (accept any variant)
        for variants in required_sections_variants:
            found = any(
                f"## {v}" in content or f"# {v}" in content
                for v in variants
            )
            if not found:
                issues.append(f"{agent_file.name}: Missing required section (one of: {', '.join(variants[:3])}...)")

        # Check recommended sections
        for section in recommended_sections:
            if f"## {section}" not in content and f"# {section}" not in content:
                warnings.append(f"{agent_file.name}: Missing recommended section '{section}'")

        # Check for pushback instruction
        if "contradictory" not in content.lower() and "pushback" not in content.lower():
            warnings.append(f"{agent_file.name}: No pushback/stop-on-ambiguity instruction found")

        # Check for write-early rule
        if "write" not in content.lower() or "early" not in content.lower():
            if "immediately" not in content.lower():
                warnings.append(f"{agent_file.name}: No write-early rule found")

    return {
        "name": "Agent Definitions",
        "passed": len(issues) == 0,
        "issues": issues,
        "warnings": warnings,
        "details": f"{len(agent_files)} agents found",
    }


def check_hooks():
    """Verify hooks are configured and scripts exist."""
    issues = []
    warnings = []

    settings_path = PROJECT_DIR / ".claude" / "settings.json"
    # Also check project-level settings
    if not settings_path.exists():
        settings_path = PROJECT_DIR / ".claude" / "settings.local.json"
    if not settings_path.exists():
        return {"name": "Hook Configuration", "passed": False,
                "issues": [".claude/settings.json not found (hooks need to be configured here)"], "warnings": []}

    settings = json.loads(settings_path.read_text())
    hooks = settings.get("hooks", {})

    if not hooks:
        issues.append("No hooks configured in settings.json")
        return {"name": "Hook Configuration", "passed": False,
                "issues": issues, "warnings": []}

    expected_events = ["SessionStart", "SessionEnd", "PreCompact", "PostToolUse"]
    for event in expected_events:
        if event not in hooks:
            warnings.append(f"No {event} hook configured")

    # Check hook scripts exist
    hooks_dir = PROJECT_DIR / ".claude" / "hooks"
    if hooks_dir.exists():
        for script in hooks_dir.glob("*.sh"):
            if not os.access(script, os.X_OK):
                issues.append(f"{script.name}: Not executable (run: chmod +x .claude/hooks/{script.name})")

    return {
        "name": "Hook Configuration",
        "passed": len(issues) == 0,
        "issues": issues,
        "warnings": warnings,
        "details": f"{len(hooks)} hook events configured",
    }


def check_learning_files():
    """Verify learning system files exist and have correct format."""
    issues = []
    warnings = []
    learning_dir = PROJECT_DIR / "agent-learnings"

    if not learning_dir.exists():
        return {"name": "Learning System", "passed": False,
                "issues": ["agent-learnings/ directory not found"], "warnings": []}

    required_files = ["PATTERNS.md", "MISTAKES.md", "DECISIONS.md", "IMPROVEMENTS.md", "METRICS.md"]
    for fname in required_files:
        fpath = learning_dir / fname
        if not fpath.exists():
            issues.append(f"Missing: agent-learnings/{fname}")
        else:
            content = fpath.read_text(encoding="utf-8", errors="ignore")
            # Check for entry format guidance
            if "Session:" not in content and "entry format" not in content.lower():
                warnings.append(f"{fname}: No entry format guidance found at top of file")

    return {
        "name": "Learning System",
        "passed": len(issues) == 0,
        "issues": issues,
        "warnings": warnings,
    }


def check_claude_md():
    """Verify CLAUDE.md exists and is appropriately sized."""
    issues = []
    warnings = []

    claude_md = PROJECT_DIR / "CLAUDE.md"
    if not claude_md.exists():
        return {"name": "CLAUDE.md", "passed": False,
                "issues": ["CLAUDE.md not found in project root"], "warnings": []}

    content = claude_md.read_text(encoding="utf-8", errors="ignore")
    line_count = content.count("\n") + 1

    if line_count > 300:
        warnings.append(f"CLAUDE.md is {line_count} lines (recommended: under 300). Consider extracting sections to reference files.")
    elif line_count > 400:
        issues.append(f"CLAUDE.md is {line_count} lines (max recommended: 300). LLMs lose instruction adherence past ~200 instructions.")

    # Check for key sections
    key_sections = ["Project Summary", "Conventions", "Pipeline", "Git"]
    for section in key_sections:
        if section.lower() not in content.lower():
            warnings.append(f"CLAUDE.md: No '{section}' section found (recommended)")

    return {
        "name": "CLAUDE.md",
        "passed": len(issues) == 0,
        "issues": issues,
        "warnings": warnings,
        "details": f"{line_count} lines",
    }


def check_python_deps():
    """Verify Python dependencies are installed."""
    issues = []
    warnings = []

    # Check sentence-transformers (required for vault-search)
    try:
        import sentence_transformers
        version = getattr(sentence_transformers, "__version__", "unknown")
        details = f"sentence-transformers {version}"
    except ImportError:
        issues.append("sentence-transformers not installed (run: pip3 install sentence-transformers)")
        details = "missing"

    # Check numpy
    try:
        import numpy
    except ImportError:
        issues.append("numpy not installed (run: pip3 install numpy)")

    return {
        "name": "Python Dependencies",
        "passed": len(issues) == 0,
        "issues": issues,
        "warnings": warnings,
        "details": details if "details" in dir() else "",
    }


def check_git():
    """Verify git is configured."""
    issues = []
    warnings = []

    try:
        result = subprocess.run(["git", "status", "--porcelain"], capture_output=True, text=True,
                                cwd=str(PROJECT_DIR), timeout=10)
        if result.returncode != 0:
            issues.append("Not a git repository (run: git init)")
    except FileNotFoundError:
        issues.append("git not found on PATH")
    except subprocess.TimeoutExpired:
        warnings.append("git status timed out")

    # Check .gitignore
    gitignore = PROJECT_DIR / ".gitignore"
    if gitignore.exists():
        content = gitignore.read_text()
        should_ignore = [".env", "node_modules", ".session", ".lab", ".codebase-map"]
        for pattern in should_ignore:
            if pattern not in content:
                warnings.append(f".gitignore: Missing '{pattern}' entry")
    else:
        warnings.append(".gitignore not found")

    return {
        "name": "Git Configuration",
        "passed": len(issues) == 0,
        "issues": issues,
        "warnings": warnings,
    }


def check_secrets():
    """Scan for accidentally committed secrets."""
    issues = []
    warnings = []

    secret_patterns = [
        (r'(?:api[_-]?key|apikey)\s*[=:]\s*["\']?[a-zA-Z0-9_\-]{20,}', "Possible API key"),
        (r'(?:password|passwd|pwd)\s*[=:]\s*["\']?[^\s"\']{8,}', "Possible password"),
        (r'sk-[a-zA-Z0-9]{20,}', "Possible OpenAI/Stripe secret key"),
        (r'[A-Z][A-Z_]{3,}_(?:KEY|SECRET|TOKEN)\s*=\s*["\']?[a-zA-Z0-9_\-]{20,}', "Possible service key"),
        (r'Bearer\s+[a-zA-Z0-9_\-\.]{20,}', "Possible Bearer token"),
    ]

    # Only scan tracked files
    try:
        result = subprocess.run(["git", "ls-files"], capture_output=True, text=True,
                                cwd=str(PROJECT_DIR), timeout=10)
        if result.returncode == 0:
            tracked_files = result.stdout.strip().split("\n")
        else:
            tracked_files = []
    except Exception:
        tracked_files = []

    # Limit scan to text files
    text_extensions = {'.md', '.txt', '.json', '.js', '.ts', '.tsx', '.jsx', '.py', '.sh', '.yml', '.yaml', '.toml'}
    skip_paths = {'node_modules', '.git', 'package-lock.json', 'yarn.lock'}

    scanned = 0
    for fname in tracked_files:
        if not fname:
            continue
        if any(skip in fname for skip in skip_paths):
            continue
        fpath = PROJECT_DIR / fname
        if not fpath.exists() or fpath.suffix not in text_extensions:
            continue
        if fpath.stat().st_size > 500_000:  # Skip large files
            continue

        try:
            content = fpath.read_text(encoding="utf-8", errors="ignore")
            scanned += 1
        except Exception:
            continue

        for pattern, label in secret_patterns:
            matches = re.findall(pattern, content)
            if matches:
                # Don't flag example/template patterns
                if any(x in fname for x in ['.example', '.template', 'GUIDE', 'README']):
                    continue
                warnings.append(f"{fname}: {label} detected ({len(matches)} match(es))")

    return {
        "name": "Secret Scan",
        "passed": len(issues) == 0,
        "issues": issues,
        "warnings": warnings,
        "details": f"Scanned {scanned} tracked files",
    }


def check_vault_search():
    """Verify vault-search is set up and index is fresh."""
    issues = []
    warnings = []

    # Check script exists
    search_script = VAULT_ROOT / "vault-search.py"
    if not search_script.exists():
        search_script = PROJECT_DIR / "vault-search.py"
    if not search_script.exists():
        issues.append("vault-search.py not found in vault root or project directory")
        return {"name": "Vault Search", "passed": False,
                "issues": issues, "warnings": []}

    # Check index exists
    index_dir = VAULT_ROOT / ".vault-search"
    if not index_dir.exists() or not (index_dir / "index.pkl").exists():
        warnings.append("Search index not built yet (run: python3 vault-search.py --reindex)")

    return {
        "name": "Vault Search",
        "passed": len(issues) == 0,
        "issues": issues,
        "warnings": warnings,
    }


def check_skills():
    """Verify skills are properly defined."""
    issues = []
    warnings = []
    skills_dir = PROJECT_DIR / ".claude" / "skills"

    if not skills_dir.exists():
        return {"name": "Skills", "passed": False,
                "issues": ["No .claude/skills/ directory"], "warnings": []}

    skill_dirs = [d for d in skills_dir.iterdir() if d.is_dir()]
    for skill_dir in skill_dirs:
        skill_file = skill_dir / "SKILL.md"
        if not skill_file.exists():
            issues.append(f"Skill '{skill_dir.name}': Missing SKILL.md")
        else:
            content = skill_file.read_text(encoding="utf-8", errors="ignore")
            if not content.startswith("---"):
                issues.append(f"Skill '{skill_dir.name}': Missing YAML frontmatter in SKILL.md")

    core_skills = ["pipeline", "startup", "closeout"]
    for skill in core_skills:
        if not (skills_dir / skill).exists():
            warnings.append(f"Core skill '{skill}' not found")

    return {
        "name": "Skills",
        "passed": len(issues) == 0,
        "issues": issues,
        "warnings": warnings,
        "details": f"{len(skill_dirs)} skills found",
    }


def check_handoffs():
    """Check handoff document health."""
    issues = []
    warnings = []
    handoffs_dir = PROJECT_DIR / "handoffs"

    if not handoffs_dir.exists():
        return {"name": "Handoff System", "passed": True,
                "issues": [], "warnings": ["No handoffs/ directory (will be created on first closeout)"]}

    handoffs = sorted(handoffs_dir.glob("HANDOFF_SESSION_*.md"))
    if not handoffs:
        warnings.append("No handoff documents found yet")
    else:
        # Check latest handoff has key sections
        latest = handoffs[-1]
        content = latest.read_text(encoding="utf-8", errors="ignore")
        key_sections = ["Context Package", "Continuation Prompt"]
        for section in key_sections:
            if section.lower() not in content.lower():
                warnings.append(f"Latest handoff ({latest.name}): Missing '{section}' section")

    return {
        "name": "Handoff System",
        "passed": len(issues) == 0,
        "issues": issues,
        "warnings": warnings,
        "details": f"{len(handoffs)} handoffs" if handoffs_dir.exists() else "",
    }


# ============================================================
# Auto-Fix
# ============================================================

def auto_fix(results):
    """Attempt to fix common issues."""
    fixes = []

    for result in results:
        for issue in result.get("issues", []):
            # Fix missing directories
            if issue.startswith("Missing directory:"):
                dir_name = issue.split(": ")[1]
                path = PROJECT_DIR / dir_name
                path.mkdir(parents=True, exist_ok=True)
                fixes.append(f"Created directory: {dir_name}")

            # Fix non-executable hooks
            if "Not executable" in issue:
                script_name = issue.split(":")[0]
                script_path = PROJECT_DIR / ".claude" / "hooks" / script_name
                if script_path.exists():
                    os.chmod(script_path, 0o755)
                    fixes.append(f"Made executable: {script_name}")

    return fixes


# ============================================================
# Output
# ============================================================

def print_results(results, quiet=False):
    """Print health check results."""
    total_issues = sum(len(r.get("issues", [])) for r in results)
    total_warnings = sum(len(r.get("warnings", [])) for r in results)
    total_passed = sum(1 for r in results if r["passed"])

    print(f"\n{C.BOLD}AI Dev Vault Health Check{C.RESET}")
    print(f"{C.DIM}{'=' * 50}{C.RESET}\n")

    for result in results:
        status = f"{C.GREEN}PASS{C.RESET}" if result["passed"] else f"{C.RED}FAIL{C.RESET}"
        details = f" {C.DIM}({result['details']}){C.RESET}" if result.get("details") else ""
        print(f"  {status}  {result['name']}{details}")

        if not quiet:
            for issue in result.get("issues", []):
                print(f"       {C.RED}! {issue}{C.RESET}")
            for warning in result.get("warnings", []):
                print(f"       {C.YELLOW}? {warning}{C.RESET}")

    print(f"\n{C.DIM}{'=' * 50}{C.RESET}")
    color = C.GREEN if total_issues == 0 else C.RED
    print(f"  {color}{total_passed}/{len(results)} checks passed{C.RESET}")
    if total_issues:
        print(f"  {C.RED}{total_issues} issue(s) to fix{C.RESET}")
    if total_warnings:
        print(f"  {C.YELLOW}{total_warnings} warning(s) to review{C.RESET}")
    print()


def print_json(results):
    """Output results as JSON."""
    output = {
        "version": __version__,
        "vault_root": str(VAULT_ROOT),
        "project_dir": str(PROJECT_DIR),
        "checks": results,
        "summary": {
            "total": len(results),
            "passed": sum(1 for r in results if r["passed"]),
            "issues": sum(len(r.get("issues", [])) for r in results),
            "warnings": sum(len(r.get("warnings", [])) for r in results),
        }
    }
    print(json.dumps(output, indent=2))


# ============================================================
# CLI
# ============================================================

def main():
    args = sys.argv[1:]

    if "--version" in args:
        print(f"health-check {__version__}")
        sys.exit(0)

    if "--help" in args:
        print(__doc__)
        sys.exit(0)

    quiet = "--quiet" in args
    fix = "--fix" in args
    as_json = "--json" in args

    # Disable colors if piping
    if not sys.stdout.isatty():
        C.disable()

    # Run all checks
    results = [
        check_directory_structure(),
        check_claude_md(),
        check_agents(),
        check_skills(),
        check_hooks(),
        check_learning_files(),
        check_vault_search(),
        check_python_deps(),
        check_git(),
        check_secrets(),
        check_handoffs(),
    ]

    if fix:
        fixes = auto_fix(results)
        if fixes:
            print(f"\n{C.BLUE}Auto-fixed:{C.RESET}")
            for f in fixes:
                print(f"  {C.GREEN}+ {f}{C.RESET}")
            print()
            # Re-run checks after fixes
            results = [
                check_directory_structure(),
                check_claude_md(),
                check_agents(),
                check_skills(),
                check_hooks(),
                check_learning_files(),
                check_vault_search(),
                check_python_deps(),
                check_git(),
                check_secrets(),
                check_handoffs(),
            ]

    if as_json:
        print_json(results)
    else:
        print_results(results, quiet)

    # Exit code
    total_issues = sum(len(r.get("issues", [])) for r in results)
    sys.exit(1 if total_issues > 0 else 0)


if __name__ == "__main__":
    main()