# -*- coding: utf-8 -*-
"""Location: ./plugins/resource_filter/resource_filter.py
Copyright 2025
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

Resource Filter Plugin - Demonstrates resource hook functionality.
This plugin demonstrates how to use resource_pre_fetch and resource_post_fetch hooks
to filter and modify resource content. It can:
- Block resources based on URI patterns or protocols
- Limit resource content size
- Redact sensitive information from resource content
- Add metadata to resources
"""

# Standard
import re
from typing import List, Pattern
from urllib.parse import urlparse

# Third-Party
from cpex.framework import (
    Plugin,
    PluginConfig,
    PluginContext,
    PluginMode,
    PluginViolation,
    ResourcePostFetchPayload,
    ResourcePostFetchResult,
    ResourcePreFetchPayload,
    ResourcePreFetchResult,
    ToolPostInvokePayload,
    ToolPostInvokeResult,
)


class ResourceFilterPlugin(Plugin):
    """Plugin that filters and modifies resources.

    This plugin demonstrates the use of resource hooks to:
    - Validate resource URIs before fetching
    - Filter content after fetching
    - Add metadata to resources
    - Block certain protocols or domains
    """

    def __init__(self, config: PluginConfig) -> None:
        """Initialize the resource filter plugin.

        Args:
            config: Plugin configuration containing filter settings.
        """
        super().__init__(config)
        plugin_config = config.config if config.config else {}
        self.max_content_size = plugin_config.get("max_content_size", 1048576)
        self.allowed_protocols = plugin_config.get("allowed_protocols", ["file", "http", "https"])
        self.blocked_domains = plugin_config.get("blocked_domains", [])
        # Precompile content filter patterns for performance
        self.content_filters: List[tuple[Pattern[str], str]] = []
        for filter_rule in plugin_config.get("content_filters", []):
            pattern = filter_rule.get("pattern")
            replacement = filter_rule.get("replacement", "***")
            if pattern:
                try:
                    compiled_pattern = re.compile(pattern, re.IGNORECASE)
                    self.content_filters.append((compiled_pattern, replacement))
                except re.error:
                    # Skip invalid patterns
                    pass

    async def resource_pre_fetch(self, payload: ResourcePreFetchPayload, context: PluginContext) -> ResourcePreFetchResult:
        """Validate and potentially modify resource requests before fetching.

        Args:
            payload: The resource pre-fetch payload containing URI and metadata.
            context: Plugin execution context.

        Returns:
            ResourcePreFetchResult indicating whether to continue and any modifications.
        """
        # Parse the URI
        try:
            parsed = urlparse(payload.uri)
        except Exception as e:
            violation = PluginViolation(reason="Invalid URI", description=f"Could not parse resource URI: {e}", code="INVALID_URI", details={"uri": payload.uri, "error": str(e)})
            return ResourcePreFetchResult(continue_processing=False, violation=violation)

        # Check if URI has a scheme
        if not parsed.scheme:
            violation = PluginViolation(reason="Invalid URI format", description="URI must have a valid scheme (protocol)", code="INVALID_URI", details={"uri": payload.uri})
            # In transform mode, log but continue
            if self.mode == PluginMode.TRANSFORM:
                return ResourcePreFetchResult(continue_processing=True, violation=violation, modified_payload=payload)
            return ResourcePreFetchResult(continue_processing=False, violation=violation)

        # Check protocol
        if parsed.scheme not in self.allowed_protocols:
            violation = PluginViolation(
                reason="Protocol not allowed",
                description=f"Protocol '{parsed.scheme}' is not in allowed list",
                code="PROTOCOL_BLOCKED",
                details={"uri": payload.uri, "protocol": parsed.scheme, "allowed": self.allowed_protocols},
            )
            # In transform mode, log but continue
            if self.mode == PluginMode.TRANSFORM:
                return ResourcePreFetchResult(continue_processing=True, violation=violation, modified_payload=payload)
            return ResourcePreFetchResult(continue_processing=False, violation=violation)

        # Check domain blocking (case-insensitive)
        if parsed.netloc:
            # Convert both to lowercase for comparison
            domain_lower = parsed.netloc.lower()
            blocked_domains_lower = [d.lower() for d in self.blocked_domains]
            if domain_lower in blocked_domains_lower or any(domain_lower.endswith("." + d) for d in blocked_domains_lower):
                violation = PluginViolation(
                    reason="Domain is blocked", description=f"Domain '{parsed.netloc}' is in blocked list", code="DOMAIN_BLOCKED", details={"uri": payload.uri, "domain": parsed.netloc}
                )
                # In transform mode, log but continue
                if self.mode == PluginMode.TRANSFORM:
                    return ResourcePreFetchResult(continue_processing=True, violation=violation, modified_payload=payload)
                return ResourcePreFetchResult(continue_processing=False, violation=violation)

        # Add metadata to track this plugin processed the request
        modified_payload = ResourcePreFetchPayload(
            uri=payload.uri,
            metadata={
                **(payload.metadata or {}),
                "validated": True,
                "protocol": parsed.scheme,
                "request_id": context.global_context.request_id,
                "user": context.global_context.user,
                "resource_filter_plugin": "pre_fetch_validated",
                "allowed_size": self.max_content_size,
            },
        )

        # Store validation info in context for post-fetch
        context.set_state("uri_validated", True)
        context.set_state("original_uri", payload.uri)

        return ResourcePreFetchResult(continue_processing=True, modified_payload=modified_payload, metadata={"validation": "passed"})

    async def resource_post_fetch(self, payload: ResourcePostFetchPayload, context: PluginContext) -> ResourcePostFetchResult:
        """Filter and modify resource content after fetching.

        Args:
            payload: The resource post-fetch payload containing fetched content.
            context: Plugin execution context.

        Returns:
            ResourcePostFetchResult with potentially modified content.
        """
        # Check if pre-fetch validation was done
        if not context.get_state("uri_validated"):
            # This resource wasn't validated in pre-fetch, skip processing
            return ResourcePostFetchResult(continue_processing=True, modified_payload=payload)

        # Process content if it's text
        modified_content = payload.content
        content_was_modified = False

        # Apply content filters if we have text content
        if hasattr(payload.content, "text") and payload.content.text:
            original_text = payload.content.text
            filtered_text = original_text

            # Check content size
            if len(filtered_text.encode("utf-8")) > self.max_content_size:
                violation = PluginViolation(
                    reason="Content exceeds maximum size",
                    description=f"Resource content exceeds maximum size of {self.max_content_size} bytes",
                    code="CONTENT_TOO_LARGE",
                    details={"uri": payload.uri, "size": len(filtered_text.encode("utf-8")), "max_size": self.max_content_size},
                )
                # In transform mode, log but continue
                if self.mode == PluginMode.TRANSFORM:
                    return ResourcePostFetchResult(continue_processing=True, violation=violation, modified_payload=payload)
                return ResourcePostFetchResult(continue_processing=False, violation=violation)

            # Apply content filters
            for compiled_pattern, replacement in self.content_filters:
                filtered_text = compiled_pattern.sub(replacement, filtered_text)

            # Update content if it was modified
            if filtered_text != original_text:
                # Create new content object with filtered text
                # First-Party
                from mcpgateway.common.models import ResourceContent

                modified_content = ResourceContent(
                    type=payload.content.type,
                    id=payload.content.id,
                    uri=payload.content.uri,
                    mime_type=getattr(payload.content, "mime_type", None),
                    text=filtered_text,
                    blob=getattr(payload.content, "blob", None),
                )
                content_was_modified = True
                context.set_state("content_filtered", True)

        # Only create modified payload if content was actually modified
        if content_was_modified:
            modified_payload = ResourcePostFetchPayload(uri=payload.uri, content=modified_content)
        else:
            # Return original payload if nothing was modified
            modified_payload = payload

        return ResourcePostFetchResult(
            continue_processing=True, modified_payload=modified_payload, metadata={"filtered": context.get_state("content_filtered", False), "original_uri": context.get_state("original_uri")}
        )

    async def tool_post_invoke(self, payload: ToolPostInvokePayload, context: PluginContext) -> ToolPostInvokeResult:
        """Handle tool invocation results.

        This plugin focuses on resource filtering, so tool invocations pass through unmodified.

        Args:
            payload: The tool invocation result payload.
            context: Plugin execution context.

        Returns:
            ToolPostInvokeResult indicating to continue processing without modifications.
        """
        # This plugin is focused on resource filtering, not tool invocations
        # Simply pass through without modification
        return ToolPostInvokeResult(continue_processing=True, modified_payload=payload)
