#!/bin/bash
#
# Governed Terraform Wrapper
# ==========================
# Enforces plan-before-apply and full audit trail.
# Part of Phase 3: Execution Pipeline.
#
# Usage:
#   tf-governed.sh plan <dir> [--target=<target>]
#   tf-governed.sh apply <dir> [--plan=<plan_file>]
#   tf-governed.sh destroy <dir> [--plan=<plan_file>]
#
# Rules:
#   1. All applies MUST have a corresponding plan artifact
#   2. Plans are stored with checksums for verification
#   3. All operations are logged to the governance ledger
#

set -euo pipefail

# Configuration
GOVERNANCE_DIR="/opt/agent-governance"
EVIDENCE_DIR="${GOVERNANCE_DIR}/evidence"
PREFLIGHT_DIR="${GOVERNANCE_DIR}/preflight"
ARTIFACT_DIR="${EVIDENCE_DIR}/terraform"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

# Ensure directories exist
mkdir -p "${ARTIFACT_DIR}"

log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_ok() {
    echo -e "${GREEN}[OK]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_block() {
    echo -e "${RED}[BLOCKED]${NC} $1"
}

# Generate unique artifact ID
generate_artifact_id() {
    echo "tf-$(date +%Y%m%d-%H%M%S)-$(openssl rand -hex 4)"
}

# Get agent info from environment or default
get_agent_info() {
    AGENT_ID="${AGENT_ID:-cli-user}"
    AGENT_TIER="${AGENT_TIER:-1}"
}

# Run preflight checks
run_preflight() {
    local targets="$1"

    log_info "Running preflight checks..."

    if [[ -f "${PREFLIGHT_DIR}/preflight.py" ]]; then
        cd "${PREFLIGHT_DIR}"
        python3 preflight.py ${targets} --action terraform --tier "${AGENT_TIER}" --agent-id "${AGENT_ID}" --quiet
        return $?
    else
        log_warn "Preflight system not available, skipping checks"
        return 0
    fi
}

# Store plan artifact
store_plan_artifact() {
    local plan_file="$1"
    local artifact_id="$2"
    local tf_dir="$3"

    local artifact_path="${ARTIFACT_DIR}/${artifact_id}"
    mkdir -p "${artifact_path}"

    # Copy plan binary
    cp "${plan_file}" "${artifact_path}/tfplan.binary"

    # Generate plan JSON
    terraform -chdir="${tf_dir}" show -json "${plan_file}" > "${artifact_path}/tfplan.json" 2>/dev/null || true

    # Generate human-readable plan
    terraform -chdir="${tf_dir}" show "${plan_file}" > "${artifact_path}/tfplan.txt" 2>/dev/null || true

    # Calculate checksums
    sha256sum "${plan_file}" > "${artifact_path}/tfplan.sha256"

    # Store metadata
    cat > "${artifact_path}/metadata.json" << EOF
{
    "artifact_id": "${artifact_id}",
    "type": "terraform_plan",
    "created_at": "$(date -u +%Y-%m-%dT%H:%M:%SZ)",
    "agent_id": "${AGENT_ID}",
    "agent_tier": ${AGENT_TIER},
    "tf_dir": "${tf_dir}",
    "checksum": "$(sha256sum ${plan_file} | cut -d' ' -f1)"
}
EOF

    echo "${artifact_path}"
}

# Verify plan artifact before apply
verify_plan_artifact() {
    local plan_file="$1"
    local artifact_id="$2"

    local artifact_path="${ARTIFACT_DIR}/${artifact_id}"

    if [[ ! -d "${artifact_path}" ]]; then
        log_error "Plan artifact not found: ${artifact_id}"
        return 1
    fi

    # Verify checksum
    local stored_checksum=$(cat "${artifact_path}/tfplan.sha256" | cut -d' ' -f1)
    local current_checksum=$(sha256sum "${plan_file}" | cut -d' ' -f1)

    if [[ "${stored_checksum}" != "${current_checksum}" ]]; then
        log_error "Plan checksum mismatch!"
        log_error "  Stored:  ${stored_checksum}"
        log_error "  Current: ${current_checksum}"
        return 1
    fi

    log_ok "Plan artifact verified: ${artifact_id}"
    return 0
}

# Log to governance ledger
log_to_ledger() {
    local action="$1"
    local status="$2"
    local details="$3"

    local ledger_entry=$(cat << EOF
{
    "timestamp": "$(date -u +%Y-%m-%dT%H:%M:%SZ)",
    "agent_id": "${AGENT_ID}",
    "agent_tier": ${AGENT_TIER},
    "tool": "terraform",
    "action": "${action}",
    "status": "${status}",
    "details": "${details}"
}
EOF
)

    echo "${ledger_entry}" >> "${EVIDENCE_DIR}/terraform-ledger.jsonl"
}

# Main command handlers
cmd_plan() {
    local tf_dir="$1"
    shift

    get_agent_info

    echo ""
    echo "=========================================="
    echo "GOVERNED TERRAFORM PLAN"
    echo "=========================================="
    echo "Agent: ${AGENT_ID} (Tier ${AGENT_TIER})"
    echo "Directory: ${tf_dir}"
    echo ""

    # Extract targets for preflight (simplified)
    local targets="sandbox-vm-01"  # Default target for preflight

    # Run preflight
    if ! run_preflight "${targets}"; then
        log_block "Preflight checks failed"
        log_to_ledger "plan" "BLOCKED" "Preflight failed"
        exit 1
    fi
    log_ok "Preflight checks passed"

    # Generate artifact ID
    local artifact_id=$(generate_artifact_id)
    local plan_file="${tf_dir}/tfplan-${artifact_id}.binary"

    log_info "Running terraform plan..."
    log_info "Plan will be stored as artifact: ${artifact_id}"

    # Run terraform plan
    if terraform -chdir="${tf_dir}" plan -out="${plan_file}" "$@"; then
        log_ok "Terraform plan completed"

        # Store artifact
        local artifact_path=$(store_plan_artifact "${plan_file}" "${artifact_id}" "${tf_dir}")
        log_ok "Plan artifact stored: ${artifact_path}"

        log_to_ledger "plan" "SUCCESS" "artifact_id=${artifact_id}"

        echo ""
        echo "=========================================="
        echo "PLAN COMPLETE"
        echo "=========================================="
        echo "Artifact ID: ${artifact_id}"
        echo "To apply this plan:"
        echo "  tf-governed.sh apply ${tf_dir} --plan=${artifact_id}"
        echo "=========================================="
    else
        log_error "Terraform plan failed"
        log_to_ledger "plan" "FAILED" "terraform error"
        exit 1
    fi
}

cmd_apply() {
    local tf_dir="$1"
    shift

    get_agent_info

    echo ""
    echo "=========================================="
    echo "GOVERNED TERRAFORM APPLY"
    echo "=========================================="
    echo "Agent: ${AGENT_ID} (Tier ${AGENT_TIER})"
    echo "Directory: ${tf_dir}"
    echo ""

    # Parse arguments
    local plan_artifact_id=""
    for arg in "$@"; do
        case $arg in
            --plan=*)
                plan_artifact_id="${arg#*=}"
                shift
                ;;
        esac
    done

    # Require plan artifact
    if [[ -z "${plan_artifact_id}" ]]; then
        log_block "Apply requires a plan artifact"
        echo ""
        echo "Usage: tf-governed.sh apply <dir> --plan=<artifact_id>"
        echo ""
        echo "First run: tf-governed.sh plan <dir>"
        echo "Then use the artifact ID provided."
        log_to_ledger "apply" "BLOCKED" "no plan artifact"
        exit 1
    fi

    local artifact_path="${ARTIFACT_DIR}/${plan_artifact_id}"
    local plan_file="${artifact_path}/tfplan.binary"

    # Verify plan exists
    if [[ ! -f "${plan_file}" ]]; then
        log_block "Plan artifact not found: ${plan_artifact_id}"
        log_to_ledger "apply" "BLOCKED" "artifact not found"
        exit 1
    fi

    # Verify plan checksum
    if ! verify_plan_artifact "${plan_file}" "${plan_artifact_id}"; then
        log_block "Plan verification failed"
        log_to_ledger "apply" "BLOCKED" "checksum mismatch"
        exit 1
    fi

    # Run preflight again before apply
    local targets="sandbox-vm-01"
    if ! run_preflight "${targets}"; then
        log_block "Preflight checks failed"
        log_to_ledger "apply" "BLOCKED" "Preflight failed"
        exit 1
    fi
    log_ok "Preflight checks passed"

    log_info "Applying plan: ${plan_artifact_id}"

    # Copy plan to tf_dir for apply
    cp "${plan_file}" "${tf_dir}/tfplan.binary"

    # Run terraform apply
    if terraform -chdir="${tf_dir}" apply "${tf_dir}/tfplan.binary"; then
        log_ok "Terraform apply completed"

        # Store apply evidence
        cat > "${artifact_path}/apply-result.json" << EOF
{
    "applied_at": "$(date -u +%Y-%m-%dT%H:%M:%SZ)",
    "agent_id": "${AGENT_ID}",
    "status": "SUCCESS"
}
EOF

        log_to_ledger "apply" "SUCCESS" "artifact_id=${plan_artifact_id}"

        echo ""
        echo "=========================================="
        echo "APPLY COMPLETE"
        echo "=========================================="
        echo "Plan artifact: ${plan_artifact_id}"
        echo "Evidence stored at: ${artifact_path}"
        echo "=========================================="
    else
        log_error "Terraform apply failed"
        log_to_ledger "apply" "FAILED" "terraform error"
        exit 1
    fi

    # Cleanup
    rm -f "${tf_dir}/tfplan.binary"
}

cmd_destroy() {
    local tf_dir="$1"
    shift

    get_agent_info

    log_warn "DESTROY requires Tier 3+ access"

    if [[ "${AGENT_TIER}" -lt 3 ]]; then
        log_block "Destroy requires Tier 3+ (current: Tier ${AGENT_TIER})"
        log_to_ledger "destroy" "BLOCKED" "insufficient tier"
        exit 1
    fi

    echo ""
    echo "=========================================="
    echo "GOVERNED TERRAFORM DESTROY"
    echo "=========================================="
    echo "Agent: ${AGENT_ID} (Tier ${AGENT_TIER})"
    echo "Directory: ${tf_dir}"
    echo ""

    log_warn "This is a destructive operation!"

    # For destroy, we still require plan-first
    log_info "To destroy, first create a destroy plan:"
    echo "  terraform -chdir=${tf_dir} plan -destroy -out=destroy.plan"
    echo "  Then use: tf-governed.sh apply ${tf_dir} --plan=<artifact_id>"

    log_to_ledger "destroy" "INFO" "destroy guidance provided"
}

# Show usage
usage() {
    echo "Governed Terraform Wrapper"
    echo ""
    echo "Usage:"
    echo "  tf-governed.sh plan <directory> [terraform options]"
    echo "  tf-governed.sh apply <directory> --plan=<artifact_id>"
    echo "  tf-governed.sh destroy <directory>"
    echo ""
    echo "Environment Variables:"
    echo "  AGENT_ID   - Agent identifier (default: cli-user)"
    echo "  AGENT_TIER - Agent trust tier 0-4 (default: 1)"
    echo ""
    echo "Examples:"
    echo "  AGENT_ID=agent-001 tf-governed.sh plan ./infra"
    echo "  tf-governed.sh apply ./infra --plan=tf-20260123-120000-abc123"
}

# Main
case "${1:-}" in
    plan)
        shift
        cmd_plan "$@"
        ;;
    apply)
        shift
        cmd_apply "$@"
        ;;
    destroy)
        shift
        cmd_destroy "$@"
        ;;
    -h|--help|help)
        usage
        ;;
    *)
        usage
        exit 1
        ;;
esac