# 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 typing import List, Optional
from uuid import UUID

from pydantic import BaseModel, Field


class ProcessingStatus(str, Enum):
    """Status of molecule processing."""

    VALID = "VALID"
    CORRUPT = "CORRUPT"


class MolecularEntity(BaseModel):
    """Representation of a processed molecular entity.

    Attributes:
        id (UUID): Unique identifier.
        source_smiles (str): Original input SMILES or string representation.
        status (ProcessingStatus): Status of processing (VALID/CORRUPT).
        canonical_smiles (Optional[str]): Canonicalized SMILES.
        selfies_string (Optional[str]): SELFIES representation.
        inchi_key (Optional[str]): The universal InChI Key hash.
        molecular_weight (Optional[float]): Molecular weight in Daltons.
        logp (Optional[float]): Partition coefficient.
        tpsa (Optional[float]): Topological Polar Surface Area.
        num_rotatable_bonds (Optional[int]): Number of rotatable bonds.
        lipinski_violations (Optional[int]): Number of Lipinski Rule of 5 violations.
        fingerprint_vector (Optional[List[int]]): Binary vector for fast search (e.g. Morgan/ECFP).
        embedding (Optional[List[float]]): Multi-modal embedding (BioCLIP).
    """

    id: UUID
    source_smiles: str
    status: ProcessingStatus

    # Fields are Optional to allow for CORRUPT state
    canonical_smiles: Optional[str] = None
    selfies_string: Optional[str] = None
    inchi_key: Optional[str] = Field(default=None, description="The universal hash")

    # Computed Properties
    molecular_weight: Optional[float] = Field(default=None, gt=0, description="Molecular weight in Daltons")
    logp: Optional[float] = Field(default=None, description="Partition coefficient")
    tpsa: Optional[float] = Field(default=None, ge=0, description="Topological Polar Surface Area")
    num_rotatable_bonds: Optional[int] = Field(default=None, ge=0, description="Number of rotatable bonds")

    # New Field
    lipinski_violations: Optional[int] = Field(
        default=None, ge=0, description="Number of Lipinski Rule of 5 violations"
    )

    fingerprint_vector: Optional[List[int]] = None  # Binary vector for fast search
    embedding: Optional[List[float]] = None  # Multi-modal embedding (BioCLIP)
