# 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 enum import Enum
from functools import partial
from pathlib import Path
from typing import Any, Dict, List, Optional, Type, Union, cast
from uuid import uuid4

import anyio
import httpx
from loguru import logger

from coreason_identity.models import UserContext
from coreason_prism.analyst import Analyst
from coreason_prism.biologist import Biologist
from coreason_prism.chemist import Chemist
from coreason_prism.embedder import Embedder
from coreason_prism.models.chemistry import MolecularEntity, ProcessingStatus
from coreason_prism.models.vision import ExtractedFigure, FigureType


class PrismMode(str, Enum):
    """Mode for processing images in Prism.

    Determines which engine handles the image request.
    """

    CHART = "CHART"
    BIO = "BIO"


class PrismAsync:
    """The Async Prism Core (AUC-1).

    Handles all logic with strict lifecycle management.
    """

    def __init__(
        self,
        chemist: Type[Chemist] = Chemist,
        analyst: Type[Analyst] = Analyst,
        biologist: Type[Biologist] = Biologist,
        embedder: Type[Embedder] = Embedder,
        light_mode: bool = False,
        client: Optional[httpx.AsyncClient] = None,
    ):
        self.chemist = chemist
        self.analyst = analyst
        self.biologist = biologist
        self.embedder = embedder
        self.light_mode = light_mode

        self._internal_client = client is None
        self._client = client or httpx.AsyncClient()

        if self.light_mode:
            logger.info("Prism initialized in LIGHT MODE. Heavy models will not be loaded.")

        # Strategy Registry for Image Processing
        self._image_handlers: Dict[PrismMode, Any] = {
            PrismMode.CHART: self.analyst.process_image,
            PrismMode.BIO: self.biologist.process_image,
        }

    async def __aenter__(self) -> "PrismAsync":
        return self

    async def __aexit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
        if self._internal_client:
            await self._client.aclose()

    async def process_molecule(self, input_data: Union[str, Path], *, context: UserContext) -> MolecularEntity:
        """Process a raw molecule string (SMILES, InChI, SELFIES) or file path.

        Delegates to Chemist, then orchestrates embedding generation via Embedder.
        """
        # 1. Detect & Transmute
        func = partial(self.chemist.process_molecule, input_data, context=context)
        entity = await anyio.to_thread.run_sync(func)
        entity = cast(MolecularEntity, entity)

        # 2. Embed (only if valid)
        if entity.status == ProcessingStatus.VALID and entity.canonical_smiles:
            # Generate BioCLIP embedding for the canonical structure
            if not self.light_mode:
                entity.embedding = await self.embed_text(entity.canonical_smiles)

        return entity

    async def validate_structure(self, structure_string: str, *, context: UserContext) -> MolecularEntity:
        """Validate a structure string (SMILES or SELFIES)."""
        func = partial(self.chemist.validate_structure, structure_string, context=context)
        return await anyio.to_thread.run_sync(func)  # type: ignore

    async def process_image(
        self,
        image_path: Union[str, Path],
        source_document_id: str,
        mode: PrismMode,
        *,
        context: UserContext,
    ) -> ExtractedFigure:
        """Process an image based on the specified mode."""
        # Validate input file existence (for both Normal and Light modes)
        path = Path(image_path)
        if not path.exists():
            logger.error(f"Image file not found: {image_path}")
            raise FileNotFoundError(f"Image file not found: {image_path}")

        if self.light_mode:
            logger.debug("Prism running in LIGHT MODE. Skipping image processing and embedding.")
            return ExtractedFigure(
                id=uuid4(),
                source_document_id=source_document_id,
                image_path=str(image_path),
                figure_type=FigureType.SKIPPED,
                data_series=None,
                embedding=None,
                metadata=None,
            )

        # 1. Dispatch to appropriate handler
        try:
            handler_mode = PrismMode(mode) if isinstance(mode, str) else mode
            handler = self._image_handlers[handler_mode]
        except (KeyError, ValueError) as e:
            raise ValueError(f"Unsupported PrismMode: {mode}") from e

        # Wrapper for handler to be run in thread
        func = partial(handler, image_path, source_document_id, context=context)
        figure = await anyio.to_thread.run_sync(func)
        figure = cast(ExtractedFigure, figure)

        # 2. Embed
        figure.embedding = await self.embed_image(image_path)

        return figure

    async def embed_text(self, text: str) -> List[float]:
        """Generate an embedding for a given text."""
        if self.light_mode:
            logger.debug("Prism is in LIGHT MODE. Text embedding skipped.")
            return []
        return await anyio.to_thread.run_sync(self.embedder.embed_text, text)  # type: ignore

    async def embed_image(self, image_path: Union[str, Path]) -> List[float]:
        """Generate an embedding for a given image."""
        if self.light_mode:
            logger.debug("Prism is in LIGHT MODE. Image embedding skipped.")
            return []
        return await anyio.to_thread.run_sync(self.embedder.embed_image, image_path)  # type: ignore


class Prism:
    """The Prism Facade (AUC-1).

    Acts as the unified integration layer/entry point for coreason-prism.
    Routes requests to the appropriate specialized engine.

    Refactored to use Dependency Injection for better testability and extensibility.
    This class now acts as a synchronous wrapper around PrismAsync.
    """

    def __init__(
        self,
        chemist: Type[Chemist] = Chemist,
        analyst: Type[Analyst] = Analyst,
        biologist: Type[Biologist] = Biologist,
        embedder: Type[Embedder] = Embedder,
        light_mode: bool = False,
        client: Optional[httpx.AsyncClient] = None,
    ):
        """Initialize Prism with injected engines.

        Defaults to the static implementations provided in the package.
        """
        self._async = PrismAsync(chemist, analyst, biologist, embedder, light_mode, client)

    def __enter__(self) -> "Prism":
        # We don't need to do much here since PrismAsync handles resources in __aenter__
        # But we should probably ensure the client is ready?
        # PrismAsync initializes client in __init__, so it's fine.
        return self

    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
        anyio.run(self._async.__aexit__, exc_type, exc_val, exc_tb)

    def process_molecule(self, input_data: Union[str, Path], *, context: UserContext) -> MolecularEntity:
        """Process a raw molecule string (SMILES, InChI, SELFIES) or file path."""

        async def wrapper() -> MolecularEntity:
            return await self._async.process_molecule(input_data, context=context)

        return anyio.run(wrapper)  # type: ignore

    def validate_structure(self, structure_string: str, *, context: UserContext) -> MolecularEntity:
        """Validate a structure string (SMILES or SELFIES)."""

        async def wrapper() -> MolecularEntity:
            return await self._async.validate_structure(structure_string, context=context)

        return anyio.run(wrapper)  # type: ignore

    def process_image(
        self,
        image_path: Union[str, Path],
        source_document_id: str,
        mode: PrismMode,
        *,
        context: UserContext,
    ) -> ExtractedFigure:
        """Process an image based on the specified mode."""

        async def wrapper() -> ExtractedFigure:
            return await self._async.process_image(image_path, source_document_id, mode, context=context)

        return anyio.run(wrapper)  # type: ignore

    def embed_text(self, text: str) -> List[float]:
        """Generate an embedding for a given text."""
        return anyio.run(self._async.embed_text, text)  # type: ignore

    def embed_image(self, image_path: Union[str, Path]) -> List[float]:
        """Generate an embedding for a given image."""
        return anyio.run(self._async.embed_image, image_path)  # type: ignore


# Default Instance for backward compatibility or simple import usage
# Usage: from coreason_prism.interface import prism
# WARNING: This instance does not support context management and may leak resources if used.
prism = Prism()
