"""
Studies endpoint module for the cBioPortal MCP server.

Contains all study-related endpoint methods:
- get_cancer_studies: List cancer studies with pagination
- search_studies: Search studies by keyword
- get_study_details: Get detailed study information
- get_multiple_studies: Fetch multiple studies concurrently
"""

import asyncio
import time
from typing import Any, Dict, List, Optional

import httpx
from .base import BaseEndpoint, handle_api_errors, validate_paginated_params
from ..utils.validation import (
    validate_page_params,
    validate_sort_params,
    validate_study_id,
    validate_keyword,
)
from ..utils.pagination import collect_all_results
from ..utils.logging import get_logger

logger = get_logger(__name__)


class StudiesEndpoints(BaseEndpoint):
    """Handles all study-related endpoints for the cBioPortal MCP server."""

    def __init__(self, api_client):
        super().__init__(api_client)


    @handle_api_errors("get cancer studies")
    @validate_paginated_params
    async def get_cancer_studies(
        self,
        page_number: int = 0,
        page_size: int = 50,
        sort_by: Optional[str] = None,
        direction: str = "ASC",
        limit: Optional[int] = None,
    ) -> Dict[str, Any]:
        """
        Get a list of cancer studies in cBioPortal with pagination support.

        Args:
            page_number: Page number to retrieve (0-based)
            page_size: Number of items per page
            sort_by: Field to sort by
            direction: Sort direction (ASC or DESC)
            limit: Maximum number of items to return across all pages (None for no limit)

        Returns:
            Dictionary containing list of studies and metadata
        """
        return await self.paginated_request(
            endpoint="studies",
            page_number=page_number,
            page_size=page_size,
            sort_by=sort_by,
            direction=direction,
            limit=limit,
            data_key="studies"
        )

    @handle_api_errors("search studies")
    async def search_studies(
        self,
        keyword: str,
        page_number: int = 0,
        page_size: int = 50,
        sort_by: Optional[str] = None,
        direction: str = "ASC",
        limit: Optional[int] = None,
    ) -> Dict[str, Any]:
        """
        Search for cancer studies by keyword in their name or description with pagination support.

        Args:
            keyword: Keyword to search for
            page_number: Page number to retrieve (0-based)
            page_size: Number of items per page
            sort_by: Field to sort by
            direction: Sort direction (ASC or DESC)
            limit: Maximum number of items to return across all pages (None for no limit)

        Returns:
            Dictionary containing list of studies and metadata
        """
        # Input Validation
        validate_keyword(keyword)
        validate_page_params(page_number, page_size, limit)
        validate_sort_params(sort_by, direction)

        try:
            # Await the asynchronous API call
            all_studies = await self.api_client.make_api_request("studies")

            if not isinstance(all_studies, list):
                # Handle cases where API request might not return a list (e.g., error response)
                error_message = (
                    "Unexpected response from API when fetching all studies."
                )
                if isinstance(all_studies, dict) and "error" in all_studies:
                    error_message = all_studies["error"]
                return {
                    "error": f"Failed to search studies for '{keyword}': {error_message}"
                }

            keyword_lower = keyword.lower()
            matching_studies = [
                study
                for study in all_studies
                if keyword_lower in study.get("name", "").lower()
                or keyword_lower in study.get("description", "").lower()
            ]

            if sort_by:
                reverse = direction.upper() == "DESC"
                matching_studies.sort(
                    key=lambda s: str(s.get(sort_by, "")), reverse=reverse
                )

            total_count = len(matching_studies)
            start_idx = page_number * page_size
            end_idx = start_idx + page_size
            paginated_studies = matching_studies[start_idx:end_idx]

            if limit and limit > 0 and len(paginated_studies) > limit:
                paginated_studies = paginated_studies[:limit]

            has_more = end_idx < total_count

            return {
                "studies": paginated_studies,
                "pagination": {
                    "page": page_number,
                    "page_size": page_size
                    if limit != 0
                    else total_count,  # Adjust page_size if all were requested
                    "total_found": total_count,
                    "has_more": has_more
                    if limit != 0
                    else False,  # No more if all were requested
                },
            }
        except httpx.HTTPStatusError as exc:
            return {
                "error": f"Failed to search studies for '{keyword}': API request failed with status {exc.response.status_code}"
            }
        except httpx.RequestError as exc:
            return {
                "error": f"Failed to search studies for '{keyword}': Network error - {str(exc)}"
            }
        except Exception as e:
            return {
                "error": f"Failed to search studies for '{keyword}': An unexpected error occurred - {str(e)}"
            }

    @handle_api_errors("get study details")
    async def get_study_details(self, study_id: str) -> Dict[str, Any]:
        """
        Get detailed information for a specific cancer study.

        Args:
            study_id: The ID of the cancer study

        Returns:
            Dictionary containing study details
        """
        # Input Validation
        validate_study_id(study_id)

        endpoint = f"studies/{study_id}"
        try:
            study = await self.api_client.make_api_request(endpoint)
            return {"study": study}
        except Exception as e:
            return {"error": f"Failed to get study details for {study_id}: {str(e)}"}

    @handle_api_errors("get multiple studies")
    async def get_multiple_studies(self, study_ids: List[str]) -> Dict:
        """
        Get details for multiple studies concurrently.

        This method demonstrates the power of async concurrency by fetching
        multiple studies in parallel, which is much faster than sequential requests.

        Args:
            study_ids: List of study IDs to fetch

        Returns:
            Dictionary mapping study IDs to their details, with metadata about the operation
        """
        if not study_ids:
            return {
                "studies": {},
                "metadata": {"count": 0, "errors": 0, "concurrent": True},
            }

        # Create a reusable async function for fetching a single study
        async def fetch_study(study_id):
            try:
                data = await self.api_client.make_api_request(f"studies/{study_id}")
                return {"study_id": study_id, "data": data, "success": True}
            except Exception as e:
                return {"study_id": study_id, "error": str(e), "success": False}

        # Create tasks for all study IDs and run them concurrently
        tasks = [fetch_study(study_id) for study_id in study_ids]
        start_time = time.perf_counter()
        # Use asyncio.gather to execute all tasks concurrently
        results = await asyncio.gather(*tasks)
        end_time = time.perf_counter()

        # Process results into a structured response
        studies_dict = {}
        error_count = 0

        for result in results:
            if result["success"]:
                studies_dict[result["study_id"]] = result["data"]
            else:
                studies_dict[result["study_id"]] = {"error": result["error"]}
                error_count += 1

        return {
            "studies": studies_dict,
            "metadata": {
                "count": len(study_ids),
                "errors": error_count,
                "execution_time": round(end_time - start_time, 3),
                "concurrent": True,
            },
        }

    @handle_api_errors("get cancer types")
    @validate_paginated_params
    async def get_cancer_types(
        self,
        page_number: int = 0,
        page_size: int = 50,
        sort_by: Optional[str] = None,
        direction: str = "ASC",
        limit: Optional[int] = None,
    ) -> Dict:
        """
        Get a list of all available cancer types in cBioPortal with pagination support.

        Args:
            page_number: Page number (0-based)
            page_size: Number of items per page
            sort_by: Field to sort by
            direction: Sort direction (ASC or DESC)
            limit: Maximum number of items to return across all pages (None for no limit)

        Returns:
            Dictionary containing list of cancer types and metadata
        """
        return await self.paginated_request(
            endpoint="cancer-types",
            page_number=page_number,
            page_size=page_size,
            sort_by=sort_by,
            direction=direction,
            limit=limit,
            data_key="cancer_types"
        )
