# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason_prism

from pathlib import Path
from typing import Any, Dict, Union
from uuid import uuid4

import torch
from loguru import logger
from PIL import Image
from transformers import SamModel, SamProcessor

from coreason_identity.models import UserContext
from coreason_prism.models.vision import ExtractedFigure, FigureType
from coreason_prism.utils.engine import HuggingFaceEngine


class Biologist(HuggingFaceEngine):
    """The Biologist (Bio-Image Engine).

    Handles segmentation and classification of medical images (e.g., Histology).
    Uses MedSAM (SAM 2 / Segment Anything) for ROI detection and cell counting.
    """

    MODEL_NAME = "flaviagiammarino/medsam-vit-base"
    MODEL_CLASS = SamModel
    PROCESSOR_CLASS = SamProcessor

    async def run(self, instruction: str, context: Dict[str, Any], *, user_context: UserContext) -> Dict[str, Any]:
        """Execute the agent task.

        Args:
            instruction (str): The task instruction.
            context (dict): The context containing data.
            user_context (UserContext): The user identity context.

        Returns:
            dict: The execution result.
        """
        logger.info(f"Biologist received instruction: {instruction}")
        image_path = context.get("image_path")
        if not image_path:
            raise ValueError("Biologist context missing 'image_path'")

        result = self.process_image(image_path, source_document_id="api_request", context=user_context)

        cell_count = result.metadata.get("cell_count") if result.metadata else 0

        return {
            "status": "success",
            "content": f"Analyzed bio-image. Cell count: {cell_count}",
            "artifacts": [{"type": "bio_analysis", "data": result.model_dump(mode="json")}],
        }

    @classmethod
    def process_image(
        cls, image_path: Union[str, Path], source_document_id: str, *, context: UserContext
    ) -> ExtractedFigure:
        """Process a bio-medical image to extract metadata (ROI, cell counts, etc.).

        Uses MedSAM to generate masks and compute metrics like tumor area percentage.

        Args:
            image_path (Union[str, Path]): Path to the image file.
            source_document_id (str): ID of the source document.
            context (UserContext): The user identity context.

        Returns:
            ExtractedFigure: A model containing the extracted metadata such as
            cell_count, tumor_area_percentage, etc.

        Raises:
            FileNotFoundError: If the image file does not exist.
            Exception: For errors during model inference.
        """
        path = Path(image_path)
        if not path.exists():
            logger.error(f"Image file not found: {path}")
            raise FileNotFoundError(f"Image file not found: {path}")

        cls._load_model()

        # Helper for type checkers
        assert cls._model is not None
        assert cls._processor is not None
        assert cls._device is not None

        try:
            logger.info("Analyzing biological sequence", user_id=context.user_id.get_secret_value())
            logger.info(f"Processing bio-image: {path}")
            image = Image.open(path).convert("RGB")

            # For MedSAM/SAM, we typically need prompts (points/boxes).
            # In "automatic" mode, we might grid sample.
            # For this implementation, we will perform a basic inference with a grid of points
            # to simulate "detecting" things.
            # A simple 2x2 grid of points.
            width, height = image.size
            input_points = [
                [[width // 3, height // 3], [2 * width // 3, 2 * height // 3]]
            ]  # Shape: (batch_size, point_batch_size, 2)

            inputs = cls._processor(image, input_points=input_points, return_tensors="pt").to(cls._device)

            with torch.no_grad():
                outputs = cls._model(**inputs)

            # outputs.pred_masks shape: (batch_size, num_masks, height, width)
            # We count the number of masks that have high confidence or just count them.
            # SAM returns 3 masks per point usually (multimask_output=True default).
            masks = outputs.pred_masks.squeeze(0)  # Remove batch dim
            iou_scores = outputs.iou_scores.squeeze(0)

            # Flatten to (N, H, W) and (N,) to simplify processing multiple points/proposals
            if masks.ndim == 4:
                # Shape is (num_points, num_proposals, H, W)
                # We flatten points and proposals into a single dimension of candidates
                N, P, H, W = masks.shape
                masks = masks.view(N * P, H, W)
                iou_scores = iou_scores.view(N * P)

            # Simple logic: Count masks with IOU > 0.8 as valid "cells/ROIs"
            valid_indices = iou_scores > 0.8
            valid_masks_count = valid_indices.sum().item()

            # Calculate Tumor Area Percentage
            tumor_area_percentage = 0.0
            if valid_masks_count > 0:
                # Extract valid masks
                selected_masks = masks[valid_indices]

                # Convert logits to binary masks (threshold > 0.0)
                binary_masks = selected_masks > 0.0

                # Compute Union of all valid masks
                # shape: (valid_count, H, W) -> logical_or across dim 0 -> (H, W)
                union_mask = binary_masks.any(dim=0)

                # Calculate percentage of pixels in the union mask
                total_pixels = union_mask.numel()
                if total_pixels > 0:
                    covered_pixels = union_mask.sum().item()
                    tumor_area_percentage = (covered_pixels / total_pixels) * 100.0

            # Metadata
            metadata: Dict[str, Any] = {
                "cell_count": int(valid_masks_count),
                "tumor_area_percentage": float(tumor_area_percentage),
                "rois_detected": int(masks.shape[0]),
                "model_used": cls.MODEL_NAME,
            }

            return ExtractedFigure(
                id=uuid4(),
                source_document_id=source_document_id,
                image_path=str(path),
                figure_type=FigureType.HISTOLOGY,
                data_series=None,
                embedding=None,  # Handled by Embedder
                metadata=metadata,
            )

        except Exception as e:
            logger.exception(f"Error processing bio-image {path}")
            raise e
