"""Audit logging: Track command execution, policy decisions, and security events.

Audit logs record:
- Command execution attempts (approved and rejected)
- Command metadata (user, host, working directory, environment)
- Policy decisions (which policy was used, which operations were allowed/denied)
- Execution results (exit code, timestamp, duration)
- Security events (policy violations, invalid syntax, SSH errors)

Design principles:
1. Comprehensive: Log all significant security events
2. Structured: Machine-readable format (JSON) for analysis
3. Safe: Never log sensitive data (passwords, keys, tokens)
4. Performant: Async logging to avoid blocking execution
5. Traceable: Each execution has unique run_id for correlation

Log levels:
- ALLOW: Command executed successfully
- DENY: Command rejected by policy
- ERROR: Command failed or execution error
- WARNING: Suspicious behavior or policy edge cases
"""

import asyncio
import json
import logging
import os
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional


@dataclass
class AuditEvent:
    """A single audit log event."""

    timestamp: float  # Unix timestamp
    run_id: int  # Unique execution identifier
    level: str  # ALLOW, DENY, ERROR, WARNING
    event_type: str  # Describe what happened

    # Command execution context
    cmd: str  # The command that was executed/rejected
    binary: str  # The primary binary being executed

    # Security context
    policy: Optional[str] = None  # Policy name used (read_only, write_capable, etc)

    # Host context
    host: Optional[str] = None  # Remote host
    username: Optional[str] = None  # SSH username
    cwd: Optional[str] = None  # Working directory

    # Results
    exit_code: Optional[int] = None  # Process exit code (-1 if timeout, None if not executed)
    duration_s: Optional[float] = None  # Execution time in seconds

    # Error information
    error: Optional[str] = None  # Error message if event_type is ERROR
    reason: Optional[str] = None  # Reason for DENY (policy violation, unsupported syntax, etc)

    # Additional structured data
    metadata: Dict = field(default_factory=dict)  # Custom data (operation counts, etc)

    def to_dict(self) -> dict:
        """Convert to dict for JSON serialization."""
        return asdict(self)

    def to_json(self) -> str:
        """Convert to JSON string."""
        return json.dumps(self.to_dict(), default=str)


class AuditLogger:
    """Audit logger for remux execution events."""

    def __init__(self, log_dir: Optional[str] = None):
        """
        Initialize audit logger.

        Args:
            log_dir: Directory to write audit logs to. If None, logs to stderr only.
        """
        self.log_dir = Path(log_dir) if log_dir else None
        self.logger = logging.getLogger("rpipe.audit")

        # Configure logger to output structured JSON
        handler = logging.StreamHandler()
        formatter = logging.Formatter("%(message)s")
        handler.setFormatter(formatter)
        self.logger.addHandler(handler)
        self.logger.setLevel(logging.DEBUG)

        # Create log file if directory specified
        self.log_file: Optional[Path] = None
        if self.log_dir:
            self.log_dir.mkdir(parents=True, exist_ok=True)
            # Daily log file: audit-YYYY-MM-DD.jsonl
            date_str = datetime.now().strftime("%Y-%m-%d")
            self.log_file = self.log_dir / f"audit-{date_str}.jsonl"

    def _maybe_rotate_log(self) -> None:
        """Check if we need to rotate to a new date's log file."""
        if not self.log_dir:
            return

        date_str = datetime.now().strftime("%Y-%m-%d")
        new_log_file = self.log_dir / f"audit-{date_str}.jsonl"

        if self.log_file != new_log_file:
            self.log_file = new_log_file

    def _write_event(self, event: AuditEvent) -> None:
        """Write event to log files."""
        # Rotate log if date changed
        self._maybe_rotate_log()

        json_str = event.to_json()

        # Write to file if configured
        if self.log_file:
            try:
                with open(self.log_file, "a", encoding="utf-8") as f:
                    f.write(json_str + "\n")
            except Exception as e:
                self.logger.error(f"Failed to write audit log: {e}")

        # Always log to stderr as fallback
        self.logger.info(json_str)

    def log_command_allowed(
        self,
        run_id: int,
        cmd: str,
        binary: str,
        host: str,
        policy: Optional[str] = None,
        username: Optional[str] = None,
        cwd: Optional[str] = None,
        metadata: Optional[Dict] = None,
    ) -> None:
        """Log that a command was allowed and executed.

        Args:
            run_id: Unique execution identifier
            cmd: Full command string
            binary: Primary binary name
            host: Remote host
            policy: Policy name used for validation
            username: SSH username
            cwd: Working directory
            metadata: Additional data (operations permitted, etc)
        """
        event = AuditEvent(
            timestamp=time.time(),
            run_id=run_id,
            level="ALLOW",
            event_type="command_executed",
            cmd=cmd,
            binary=binary,
            policy=policy,
            host=host,
            username=username,
            cwd=cwd,
            metadata=metadata or {},
        )
        self._write_event(event)

    def log_command_completed(
        self,
        run_id: int,
        exit_code: int,
        duration_s: float,
    ) -> None:
        """Log command completion with results.

        Args:
            run_id: Unique execution identifier
            exit_code: Process exit code
            duration_s: Execution time in seconds
        """
        event = AuditEvent(
            timestamp=time.time(),
            run_id=run_id,
            level="ALLOW",
            event_type="command_completed",
            cmd="",  # Inferred from earlier log
            binary="",
            exit_code=exit_code,
            duration_s=duration_s,
        )
        self._write_event(event)

    def log_command_denied(
        self,
        run_id: int,
        cmd: str,
        binary: str,
        reason: str,
        policy: Optional[str] = None,
        host: Optional[str] = None,
        username: Optional[str] = None,
    ) -> None:
        """Log that a command was rejected by policy.

        Args:
            run_id: Unique execution identifier
            cmd: Full command string
            binary: Primary binary name
            reason: Why the command was denied (policy violation, etc)
            policy: Policy name that rejected the command
            host: Remote host (if available)
            username: SSH username (if available)
        """
        event = AuditEvent(
            timestamp=time.time(),
            run_id=run_id,
            level="DENY",
            event_type="command_rejected",
            cmd=cmd,
            binary=binary,
            reason=reason,
            policy=policy,
            host=host,
            username=username,
        )
        self._write_event(event)

    def log_command_error(
        self,
        run_id: int,
        cmd: str,
        binary: str,
        error: str,
        error_type: str = "execution_error",
        host: Optional[str] = None,
        username: Optional[str] = None,
    ) -> None:
        """Log an error during command execution.

        Args:
            run_id: Unique execution identifier
            cmd: Full command string
            binary: Primary binary name
            error: Error message
            error_type: Type of error (execution_error, ssh_error, timeout, etc)
            host: Remote host
            username: SSH username
        """
        event = AuditEvent(
            timestamp=time.time(),
            run_id=run_id,
            level="ERROR",
            event_type=error_type,
            cmd=cmd,
            binary=binary,
            error=error,
            host=host,
            username=username,
        )
        self._write_event(event)

    def log_syntax_error(
        self,
        run_id: int,
        cmd: str,
        error: str,
    ) -> None:
        """Log a syntax or parsing error.

        Args:
            run_id: Unique execution identifier
            cmd: Full command string
            error: Error message
        """
        event = AuditEvent(
            timestamp=time.time(),
            run_id=run_id,
            level="ERROR",
            event_type="syntax_error",
            cmd=cmd,
            binary="",
            error=error,
        )
        self._write_event(event)

    def log_security_event(
        self,
        run_id: int,
        event_type: str,
        message: str,
        cmd: Optional[str] = None,
        host: Optional[str] = None,
    ) -> None:
        """Log a security event (suspicious behavior, etc).

        Args:
            run_id: Unique execution identifier
            event_type: Type of security event
            message: Description of the event
            cmd: Command involved (if applicable)
            host: Host involved (if applicable)
        """
        event = AuditEvent(
            timestamp=time.time(),
            run_id=run_id,
            level="WARNING",
            event_type=event_type,
            cmd=cmd or "",
            binary="",
            reason=message,
            host=host,
        )
        self._write_event(event)


# Global audit logger instance
_audit_logger: Optional[AuditLogger] = None


def init_audit_logger(log_dir: Optional[str] = None) -> AuditLogger:
    """Initialize the global audit logger.

    Args:
        log_dir: Directory to write audit logs to

    Returns:
        AuditLogger instance
    """
    global _audit_logger
    _audit_logger = AuditLogger(log_dir)
    return _audit_logger


def get_audit_logger() -> AuditLogger:
    """Get the global audit logger (creates one if needed)."""
    global _audit_logger
    if _audit_logger is None:
        _audit_logger = AuditLogger()
    return _audit_logger
