"""Lambda handler for policy evaluation.

This module implements the AWS Lambda function for evaluating code against
MergeGuide security policies. It serves as the core evaluation endpoint for
the MergeGuide policy enforcement system.

Triggers:
    - API Gateway HTTP API: POST /evaluate
    - Direct Lambda invocation from other AWS services
    - Step Functions workflow integration

Event Sources:
    1. API Gateway (HTTP API v2 format):
        {
            "body": "{\"code\": \"...\", \"file_path\": \"...\", \"context\": {...}}",
            "headers": {...},
            "requestContext": {...}
        }

    2. Direct Invocation:
        {
            "code": "...",
            "file_path": "...",
            "context": {...}
        }

Response Format:
    {
        "statusCode": 200,
        "body": {
            "evaluation_id": "uuid",
            "passed": true/false,
            "summary": {...},
            "violations": [...],
            "evidence": {...},
            "duration_ms": 123
        },
        "headers": {...}
    }

Environment Variables:
    ENVIRONMENT: Deployment environment (dev, staging, prod)
    POLICIES_TABLE: DynamoDB table for custom policies
    EVALUATIONS_TABLE: DynamoDB table for evaluation history
    EVIDENCE_BUCKET: S3 bucket for evidence artifacts
    SIGNING_SECRET_ARN: ARN for evidence signing key in Secrets Manager
    AWS_REGION: AWS region for service clients

Architecture:
    This Lambda uses lazy initialization for expensive resources (PolicyEngine,
    EvidenceStorage, signing keys) to optimize cold start times. Resources are
    cached in module-level globals and reused across warm invocations.
"""

from __future__ import annotations

import json
import logging
import os
from typing import Any

# Configure logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# Import mergeguide modules - these will be bundled with the Lambda
try:
    from mergeguide.engine.evaluator import (
        EvaluationContext,
        InputType,
        PolicyEngine,
        PolicyInput,
    )
    from mergeguide.evidence.artifact import EvidenceArtifact, SourceInfo
    from mergeguide.evidence.storage import EvidenceStorage
    from mergeguide.policies.defaults import get_default_rules
    from mergeguide.policies.loader import DynamoDBPolicyLoader
except ImportError:
    # Running locally or in test mode
    import sys

    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    from mergeguide.engine.evaluator import (
        EvaluationContext,
        InputType,
        PolicyEngine,
        PolicyInput,
    )
    from mergeguide.evidence.artifact import EvidenceArtifact, SourceInfo
    from mergeguide.evidence.storage import EvidenceStorage
    from mergeguide.policies.defaults import get_default_rules
    from mergeguide.policies.loader import DynamoDBPolicyLoader


# Environment variables
ENVIRONMENT = os.environ.get("ENVIRONMENT", "dev")
POLICIES_TABLE = os.environ.get("POLICIES_TABLE", f"sentinel-policies-{ENVIRONMENT}")
EVALUATIONS_TABLE = os.environ.get("EVALUATIONS_TABLE", f"sentinel-evaluations-{ENVIRONMENT}")
EVIDENCE_BUCKET = os.environ.get("EVIDENCE_BUCKET", f"sentinel-evidence-{ENVIRONMENT}")
SIGNING_SECRET_ARN = os.environ.get("SIGNING_SECRET_ARN", "")
AWS_REGION = os.environ.get("AWS_REGION", "us-east-1")

# Lazy-loaded resources
_engine: PolicyEngine | None = None
_storage: EvidenceStorage | None = None
_signing_key: str | None = None


def get_engine() -> PolicyEngine:
    """Get or create the singleton PolicyEngine instance.

    Uses lazy initialization to defer expensive PolicyEngine setup until
    first use. The engine is cached in a module-level global to persist
    across warm Lambda invocations, significantly reducing latency.

    The engine is initialized with default security rules (secrets detection,
    dangerous function detection, etc.) and optionally loads organization-specific
    rules from DynamoDB.

    Returns:
        PolicyEngine: Configured policy engine ready for evaluations.

    Note:
        DynamoDB policy loading failures are logged but do not prevent
        the engine from functioning with default rules.
    """
    global _engine
    if _engine is None:
        # Initialize with built-in default security rules
        _engine = PolicyEngine(rules=get_default_rules())

        # Attempt to load organization-specific rules from DynamoDB
        # This enables per-org policy customization
        try:
            DynamoDBPolicyLoader(
                table_name=POLICIES_TABLE,
                region=AWS_REGION,
            )
            # Note: org_id would come from the request context
            # For now, we just use defaults
        except Exception as e:
            logger.warning(f"Failed to initialize DynamoDB loader: {e}")

    return _engine


def get_storage() -> EvidenceStorage:
    """Get or create the singleton EvidenceStorage instance.

    Uses lazy initialization to defer S3 client setup until first use.
    The storage instance is cached to reuse the underlying boto3 client
    across warm invocations.

    Returns:
        EvidenceStorage: Configured storage client for the evidence S3 bucket.
    """
    global _storage
    if _storage is None:
        _storage = EvidenceStorage(
            bucket_name=EVIDENCE_BUCKET,
            region=AWS_REGION,
        )
    return _storage


def get_signing_key() -> str:
    """Retrieve the HMAC signing key from AWS Secrets Manager.

    The signing key is used to create cryptographic signatures on evidence
    artifacts, enabling verification of artifact integrity and authenticity.
    The key is fetched once and cached for subsequent invocations.

    Returns:
        str: The signing key string, or empty string if unavailable.

    Note:
        - Returns empty string if SIGNING_SECRET_ARN is not configured
        - Failures are logged as warnings but do not raise exceptions
        - Evidence artifacts created without a signing key will be unsigned
    """
    global _signing_key
    if _signing_key is None and SIGNING_SECRET_ARN:
        try:
            import boto3

            # Fetch the secret containing the signing key
            client = boto3.client("secretsmanager", region_name=AWS_REGION)
            response = client.get_secret_value(SecretId=SIGNING_SECRET_ARN)
            secret = json.loads(response["SecretString"])
            _signing_key = secret.get("signing_key", "")
        except Exception as e:
            logger.warning(f"Failed to get signing key: {e}")
            _signing_key = ""
    return _signing_key or ""


def handler(event: dict[str, Any], context: Any) -> dict[str, Any]:
    """AWS Lambda entry point for policy evaluation requests.

    This is the main handler function invoked by AWS Lambda. It orchestrates
    the complete evaluation workflow:
    1. Parse and validate the incoming request
    2. Run policy evaluation against the provided code
    3. Generate and store a signed evidence artifact
    4. Return the evaluation results

    Supports two invocation patterns:
        - API Gateway: Event contains "body" key with JSON string payload
        - Direct invocation: Event IS the payload (no "body" wrapper)

    Args:
        event: Lambda event dictionary. For API Gateway invocations, contains:
            - body (str): JSON-encoded request payload
            - headers (dict): HTTP headers
            - requestContext (dict): API Gateway context
            For direct invocations, the event IS the request payload.
        context: AWS Lambda context object providing runtime information
            (request ID, function name, memory limit, etc.). Not currently used
            but required by the Lambda handler signature.

    Returns:
        dict: API Gateway-compatible response with:
            - statusCode (int): HTTP status code (200, 400, or 500)
            - body (str): JSON-encoded response payload
            - headers (dict): Response headers including Content-Type

    Raises:
        No exceptions are raised; all errors are caught and returned as
        appropriate HTTP error responses (400 for validation errors,
        500 for internal errors).

    Example:
        Direct invocation payload::

            {
                "code": "password = 'secret123'",
                "file_path": "config.py",
                "context": {
                    "org_id": "acme-corp",
                    "repo": "backend-api",
                    "branch": "feature/auth"
                }
            }

    Response payload (in body)::

        {
            "evaluation_id": "eval-abc123",
            "passed": false,
            "summary": {"total_violations": 1, "errors": 1, ...},
            "violations": [...],
            "evidence": {"artifact_id": "...", "storage_key": "..."},
            "duration_ms": 45
        }
    """
    # Log truncated event for debugging (avoid logging sensitive code content)
    logger.info(f"Received event: {json.dumps(event)[:500]}")

    try:
        # Step 1: Parse the request body from API Gateway or direct invocation format
        body = _parse_body(event)

        # Step 2: Build the policy input from the request
        # This determines whether we're evaluating code, diff, or raw content
        input_data = _build_input(body)

        # Step 3: Build the evaluation context (org, repo, branch, etc.)
        eval_context = _build_context(body)

        # Step 4: Run the policy evaluation
        engine = get_engine()
        output = engine.evaluate(input_data)

        # Step 5: Generate evidence artifact for audit trail
        # The artifact captures the complete evaluation record
        source = SourceInfo(
            type=input_data.input_type.value,
            location=input_data.file_path or "inline",
            ref=eval_context.commit_sha or eval_context.branch,
            size_bytes=len(input_data.content),
        )

        artifact = EvidenceArtifact.from_output(
            output=output,
            source=source,
            context=eval_context,
            input_content=input_data.content,
            policies_used=[r.id for r in engine.rule_engine.rules],
        )

        # Step 6: Sign the artifact for integrity verification
        signing_key = get_signing_key()
        if signing_key:
            artifact.sign(signing_key)

        # Step 7: Store evidence in S3 for compliance and audit purposes
        storage = get_storage()
        try:
            evidence_key = storage.store(
                artifact=artifact,
                org_id=eval_context.org_id,
                repo=eval_context.repo or "unknown",
            )
            logger.info(f"Stored evidence at: {evidence_key}")
        except Exception as e:
            # Evidence storage failure is non-fatal; evaluation still succeeds
            logger.warning(f"Failed to store evidence: {e}")
            evidence_key = None

        # Step 8: Build the API Gateway-compatible response
        response = {
            "statusCode": 200,
            "body": json.dumps({
                "evaluation_id": output.evaluation_id,
                "passed": output.passed,
                "summary": {
                    "total_violations": len(output.violations),
                    "errors": sum(1 for v in output.violations if v.severity == "error"),
                    "warnings": sum(1 for v in output.violations if v.severity == "warning"),
                    "infos": sum(1 for v in output.violations if v.severity == "info"),
                    "files_checked": output.files_checked,
                },
                "violations": [
                    {
                        "rule_id": v.rule_id,
                        "rule_name": v.rule_name,
                        "message": v.message,
                        "severity": v.severity,
                        "file_path": v.file_path,
                        "line_no": v.line_no,
                        "column_start": v.column_start,
                        "column_end": v.column_end,
                        "remediation": v.remediation,
                    }
                    for v in output.violations
                ],
                "evidence": {
                    "artifact_id": artifact.id,
                    "storage_key": evidence_key,
                } if evidence_key else None,
                "duration_ms": output.duration_ms,
            }),
            "headers": {
                "Content-Type": "application/json",
                # Custom header for easy access to evaluation ID without parsing body
                "X-MergeGuide-Evaluation-Id": output.evaluation_id,
            },
        }

        return response

    except ValueError as e:
        # Validation errors (missing required fields, invalid input type)
        logger.error(f"Validation error: {e}")
        return {
            "statusCode": 400,
            "body": json.dumps({"error": str(e)}),
            "headers": {"Content-Type": "application/json"},
        }

    except Exception as e:
        # Unexpected errors - log full stack trace for debugging
        logger.exception(f"Evaluation error: {e}")
        return {
            "statusCode": 500,
            "body": json.dumps({"error": "Internal evaluation error"}),
            "headers": {"Content-Type": "application/json"},
        }


def _parse_body(event: dict[str, Any]) -> dict[str, Any]:
    """Parse and extract the request body from various Lambda event formats.

    Handles the differences between API Gateway invocations (where the payload
    is wrapped in a "body" key as a JSON string) and direct Lambda invocations
    (where the event IS the payload).

    Args:
        event: The raw Lambda event dictionary.

    Returns:
        dict: The parsed request body as a dictionary.

    Raises:
        json.JSONDecodeError: If body is a string but not valid JSON.
    """
    # API Gateway HTTP API v2 format wraps the payload in "body"
    if "body" in event:
        body = event["body"]
        # API Gateway may pass body as JSON string or already-parsed dict
        if isinstance(body, str):
            return json.loads(body)
        return body

    # Direct invocation format - the event itself is the payload
    return event


def _build_input(body: dict[str, Any]) -> PolicyInput:
    """Construct a PolicyInput from the request body.

    Determines the input type based on which field is present in the request
    and creates the appropriate PolicyInput object for evaluation.

    Input Types:
        - "diff": Git diff content for PR-based evaluation
        - "code": Source code file content (most common)
        - "content": Generic content for flexible use cases

    Args:
        body: Parsed request body dictionary containing one of:
            - diff (str): Git diff content
            - code (str): Source code content
            - content (str): Generic content
            And optionally:
            - file_path (str): Path to the file being evaluated

    Returns:
        PolicyInput: Input object ready for policy evaluation.

    Raises:
        ValueError: If none of the required content fields are present.
    """
    # Check for each input type in order of specificity
    if "diff" in body:
        # Diff input for PR evaluations - file_path not applicable
        return PolicyInput(
            input_type=InputType.DIFF,
            content=body["diff"],
            file_path="",
        )
    elif "code" in body:
        # File input - the most common case for single-file evaluation
        return PolicyInput(
            input_type=InputType.FILE,
            content=body["code"],
            file_path=body.get("file_path", ""),
        )
    elif "content" in body:
        # Generic content input for flexible use cases
        return PolicyInput(
            input_type=InputType.CONTENT,
            content=body["content"],
            file_path=body.get("file_path", ""),
        )
    else:
        raise ValueError("Request must include 'code', 'diff', or 'content' field")


def _build_context(body: dict[str, Any]) -> EvaluationContext:
    """Construct an EvaluationContext from the request body.

    Extracts optional context information that enriches the evaluation with
    metadata about the source of the code being evaluated. This context is
    used for:
    - Organization-specific policy loading
    - Evidence artifact storage paths
    - Audit trail information

    Args:
        body: Parsed request body dictionary optionally containing:
            context (dict): Nested context object with:
                - org_id (str): Organization identifier (default: "default")
                - repo (str): Repository name (e.g., "owner/repo")
                - branch (str): Git branch name
                - commit_sha (str): Full commit SHA
                - author (str): Commit or PR author
                - pr_number (int): Pull request number if applicable
                - metadata (dict): Arbitrary key-value metadata

    Returns:
        EvaluationContext: Context object with all available metadata.
    """
    # Extract the nested context object, defaulting to empty dict
    context_data = body.get("context", {})

    return EvaluationContext(
        org_id=context_data.get("org_id", "default"),
        repo=context_data.get("repo", ""),
        branch=context_data.get("branch", ""),
        commit_sha=context_data.get("commit_sha", ""),
        author=context_data.get("author", ""),
        pr_number=context_data.get("pr_number"),
        metadata=context_data.get("metadata", {}),
    )


# For local testing
if __name__ == "__main__":
    # Test evaluation
    test_event = {
        "code": """
API_KEY = "sk-1234567890abcdef1234567890abcdef"
password = "super_secret_password"

def process():
    print("Processing...")
    # TODO: Add error handling
    eval("some_code")
""",
        "file_path": "test.py",
        "context": {
            "org_id": "test-org",
            "repo": "test-repo",
            "branch": "main",
        },
    }

    result = handler(test_event, None)
    print(json.dumps(json.loads(result["body"]), indent=2))
