# -*- coding: utf-8 -*-
"""Location: ./mcpgateway/services/completion_service.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

Completion Service Implementation.
This module implements argument completion according to the MCP specification.
It handles completion suggestions for prompt arguments and resource URIs.

Examples:
    >>> from mcpgateway.services.completion_service import CompletionService, CompletionError
    >>> service = CompletionService()
    >>> isinstance(service, CompletionService)
    True
    >>> service._custom_completions
    {}
"""

# Standard
from typing import Any, Dict, List, Optional

# Third-Party
from sqlalchemy import desc, select
from sqlalchemy.orm import Session

# First-Party
from mcpgateway.common.models import CompleteResult
from mcpgateway.db import Prompt as DbPrompt
from mcpgateway.db import Resource as DbResource
from mcpgateway.services.logging_service import LoggingService

# Initialize logging service first
logging_service = LoggingService()
logger = logging_service.get_logger(__name__)


class CompletionError(Exception):
    """Base class for completion errors.

    Examples:
        >>> from mcpgateway.services.completion_service import CompletionError
        >>> err = CompletionError("Invalid reference")
        >>> str(err)
        'Invalid reference'
        >>> isinstance(err, Exception)
        True
    """


class CompletionService:
    """MCP completion service.

    Handles argument completion for:
    - Prompt arguments based on schema
    - Resource URIs with templates
    - Custom completion sources
    """

    def __init__(self):
        """Initialize completion service.

        Examples:
            >>> from mcpgateway.services.completion_service import CompletionService
            >>> service = CompletionService()
            >>> hasattr(service, '_custom_completions')
            True
            >>> service._custom_completions
            {}
        """
        self._custom_completions: Dict[str, List[str]] = {}

    async def initialize(self) -> None:
        """Initialize completion service."""
        logger.info("Initializing completion service")

    async def shutdown(self) -> None:
        """Shutdown completion service."""
        logger.info("Shutting down completion service")
        self._custom_completions.clear()

    async def handle_completion(
        self,
        db: Session,
        request: Dict[str, Any],
        user_email: Optional[str] = None,
        token_teams: Optional[List[str]] = None,
    ) -> CompleteResult:
        """Handle completion request.

        Args:
            db: Database session
            request: Completion request
            user_email: Caller email used for owner/team visibility checks
            token_teams: Normalized token teams (`None` admin bypass, `[]` public-only, list for team scope)

        Returns:
            Completion result with suggestions

        Raises:
            CompletionError: If completion fails

        Examples:
            >>> from mcpgateway.services.completion_service import CompletionService
            >>> from unittest.mock import MagicMock
            >>> service = CompletionService()
            >>> db = MagicMock()
            >>> request = {'ref': {'type': 'ref/prompt', 'name': 'prompt1'}, 'argument': {'name': 'arg1', 'value': ''}}
            >>> db.execute.return_value.scalars.return_value.all.return_value = []
            >>> import asyncio
            >>> try:
            ...     asyncio.run(service.handle_completion(db, request))
            ... except Exception:
            ...     pass
        """
        try:
            # Get reference and argument info
            ref = request.get("ref", {})
            ref_type = ref.get("type")
            arg = request.get("argument", {})
            arg_name = arg.get("name")
            arg_value = arg.get("value", "")

            if not ref_type or not arg_name:
                raise CompletionError("Missing reference type or argument name")

            # Handle different reference types
            if ref_type == "ref/prompt":
                result = await self._complete_prompt_argument(db, ref, arg_name, arg_value, user_email=user_email, token_teams=token_teams)
            elif ref_type == "ref/resource":
                result = await self._complete_resource_uri(db, ref, arg_value, user_email=user_email, token_teams=token_teams)
            else:
                raise CompletionError(f"Invalid reference type: {ref_type}")

            return result

        except Exception as e:
            logger.error(f"Completion error: {e}")
            raise CompletionError(str(e))

    async def _resolve_team_ids(self, db: Session, user_email: Optional[str], token_teams: Optional[List[str]]) -> List[str]:
        """Resolve effective team IDs for scoped visibility checks.

        Args:
            db: Database session
            user_email: Caller email for DB-based team lookup when token teams are not explicit
            token_teams: Explicit token team scope when present

        Returns:
            Effective team IDs used to build visibility filters.
        """
        if token_teams is not None:
            return token_teams
        if not user_email:
            return []

        # First-Party
        from mcpgateway.services.team_management_service import TeamManagementService  # pylint: disable=import-outside-toplevel

        team_service = TeamManagementService(db)
        user_teams = await team_service.get_user_teams(user_email)
        return [team.id for team in user_teams]

    @staticmethod
    def _apply_visibility_scope(stmt, model, user_email: Optional[str], token_teams: Optional[List[str]], team_ids: List[str], db: Session):
        """Thin passthrough to :meth:`BaseService._apply_visibility_scope`.

        Kept as a method on this class for API stability with existing
        callers; the actual logic (including the admin-bypass caller
        contract) lives in :class:`BaseService`.

        Args:
            stmt: SQLAlchemy statement to constrain
            model: ORM model with visibility/team/owner columns
            user_email: Caller email used for owner visibility
            token_teams: Explicit token team scope when present
            team_ids: Effective team IDs for team visibility
            db: Required session for the admin bypass check.

        Returns:
            Scoped SQLAlchemy statement.
        """
        # First-Party
        from mcpgateway.services.base_service import BaseService  # pylint: disable=import-outside-toplevel

        return BaseService._apply_visibility_scope(stmt, model, user_email, token_teams, team_ids, db)  # pylint: disable=protected-access

    async def _complete_prompt_argument(
        self,
        db: Session,
        ref: Dict[str, Any],
        arg_name: str,
        arg_value: str,
        user_email: Optional[str] = None,
        token_teams: Optional[List[str]] = None,
    ) -> CompleteResult:
        """Complete prompt argument value.

        Args:
            db: Database session
            ref: Prompt reference
            arg_name: Argument name
            arg_value: Current argument value
            user_email: Caller email used for owner/team visibility checks
            token_teams: Normalized token teams (`None` admin bypass, `[]` public-only, list for team scope)

        Returns:
            Completion suggestions

        Raises:
            CompletionError: If prompt is missing or not found

        Examples:
            >>> from mcpgateway.services.completion_service import CompletionService, CompletionError
            >>> from unittest.mock import MagicMock
            >>> import asyncio
            >>> service = CompletionService()
            >>> db = MagicMock()

            >>> # Test missing prompt name
            >>> ref = {}
            >>> try:
            ...     asyncio.run(service._complete_prompt_argument(db, ref, 'arg1', 'val'))
            ... except CompletionError as e:
            ...     str(e)
            'Missing prompt name'

            >>> # Test custom completions
            >>> service.register_completions('color', ['red', 'green', 'blue'])
            >>> db.execute.return_value.scalar_one_or_none.return_value = MagicMock(
            ...     argument_schema={'properties': {'color': {'name': 'color'}}}
            ... )
            >>> result = asyncio.run(service._complete_prompt_argument(
            ...     db, {'name': 'test'}, 'color', 'r'
            ... ))
            >>> result.completion['values']
            ['red', 'green']
        """
        # Get prompt
        prompt_name = ref.get("name")
        if not prompt_name:
            raise CompletionError("Missing prompt name")

        # Only consider prompts that are enabled and visible to caller
        team_ids = await self._resolve_team_ids(db, user_email, token_teams)
        stmt = select(DbPrompt).where(DbPrompt.name == prompt_name).where(DbPrompt.enabled)  # pylint: disable=comparison-with-callable
        stmt = self._apply_visibility_scope(stmt, DbPrompt, user_email=user_email, token_teams=token_teams, team_ids=team_ids, db=db)
        stmt = stmt.order_by(desc(DbPrompt.created_at), desc(DbPrompt.id)).limit(1)
        prompt = db.execute(stmt).scalar_one_or_none()

        if not prompt:
            raise CompletionError(f"Prompt not found: {prompt_name}")

        # Find argument in schema
        arg_schema = None
        for arg in prompt.argument_schema.get("properties", {}).values():
            if arg.get("name") == arg_name:
                arg_schema = arg
                break

        if not arg_schema:
            raise CompletionError(f"Argument not found: {arg_name}")

        # Get enum values if defined
        if "enum" in arg_schema:
            values = [v for v in arg_schema["enum"] if arg_value.lower() in str(v).lower()]
            return CompleteResult(
                completion={
                    "values": values[:100],
                    "total": len(values),
                    "hasMore": len(values) > 100,
                }
            )

        # Check custom completions
        if arg_name in self._custom_completions:
            values = [v for v in self._custom_completions[arg_name] if arg_value.lower() in v.lower()]
            return CompleteResult(
                completion={
                    "values": values[:100],
                    "total": len(values),
                    "hasMore": len(values) > 100,
                }
            )

        # No completions available
        return CompleteResult(completion={"values": [], "total": 0, "hasMore": False})

    async def _complete_resource_uri(
        self,
        db: Session,
        ref: Dict[str, Any],
        arg_value: str,
        user_email: Optional[str] = None,
        token_teams: Optional[List[str]] = None,
    ) -> CompleteResult:
        """Complete resource URI.

        Args:
            db: Database session
            ref: Resource reference
            arg_value: Current URI value
            user_email: Caller email used for owner/team visibility checks
            token_teams: Normalized token teams (`None` admin bypass, `[]` public-only, list for team scope)

        Returns:
            URI completion suggestions

        Raises:
            CompletionError: If URI template is missing

        Examples:
            >>> from mcpgateway.services.completion_service import CompletionService, CompletionError
            >>> from unittest.mock import MagicMock
            >>> import asyncio
            >>> service = CompletionService()
            >>> db = MagicMock()

            >>> # Test missing URI template
            >>> ref = {}
            >>> try:
            ...     asyncio.run(service._complete_resource_uri(db, ref, 'test'))
            ... except CompletionError as e:
            ...     str(e)
            'Missing URI template'

            >>> # Test resource filtering
            >>> ref = {'uri': 'template://'}
            >>> mock_resources = [
            ...     MagicMock(uri='file://doc1.txt'),
            ...     MagicMock(uri='file://doc2.txt'),
            ...     MagicMock(uri='http://example.com')
            ... ]
            >>> db.execute.return_value.scalars.return_value.all.return_value = mock_resources
            >>> result = asyncio.run(service._complete_resource_uri(db, ref, 'doc'))
            >>> len(result.completion['values'])
            2
            >>> 'file://doc1.txt' in result.completion['values']
            True
        """
        # Get base URI template
        uri_template = ref.get("uri")
        if not uri_template:
            raise CompletionError("Missing URI template")

        # List matching resources visible to caller
        team_ids = await self._resolve_team_ids(db, user_email, token_teams)
        stmt = select(DbResource).where(DbResource.enabled)
        stmt = self._apply_visibility_scope(stmt, DbResource, user_email=user_email, token_teams=token_teams, team_ids=team_ids, db=db)
        resources = db.execute(stmt).scalars().all()

        # Filter by URI pattern
        matches = []
        for resource in resources:
            if arg_value.lower() in resource.uri.lower():
                matches.append(resource.uri)

        return CompleteResult(
            completion={
                "values": matches[:100],
                "total": len(matches),
                "hasMore": len(matches) > 100,
            }
        )

    def register_completions(self, arg_name: str, values: List[str]) -> None:
        """Register custom completion values.

        Args:
            arg_name: Argument name
            values: Completion values

        Examples:
            >>> from mcpgateway.services.completion_service import CompletionService
            >>> service = CompletionService()
            >>> service.register_completions('arg1', ['a', 'b'])
            >>> service._custom_completions['arg1']
            ['a', 'b']
            >>> service.register_completions('arg2', ['x', 'y', 'z'])
            >>> len(service._custom_completions)
            2
            >>> service.register_completions('arg1', ['c'])  # Overwrite
            >>> service._custom_completions['arg1']
            ['c']
        """
        self._custom_completions[arg_name] = list(values)

    def unregister_completions(self, arg_name: str) -> None:
        """Unregister custom completion values.

        Args:
            arg_name: Argument name

        Examples:
            >>> from mcpgateway.services.completion_service import CompletionService
            >>> service = CompletionService()
            >>> service.register_completions('arg1', ['a', 'b'])
            >>> service.unregister_completions('arg1')
            >>> 'arg1' in service._custom_completions
            False
        """
        self._custom_completions.pop(arg_name, None)
