# 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, List, Union
from uuid import uuid4

import datamol as dm
import selfies as sf
from loguru import logger
from rdkit import Chem

from coreason_identity.models import UserContext
from coreason_prism.models.chemistry import MolecularEntity, ProcessingStatus


class Chemist:
    """The Chemist (Molecular Engine).

    Handles molecular normalization, transmutation (SMILES -> SELFIES),
    property calculation, and fingerprint generation.
    """

    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 (e.g., smiles).
            user_context (UserContext): The user identity context.

        Returns:
            dict: The execution result.
        """
        logger.info(f"Chemist received instruction: {instruction}")
        input_data = context.get("smiles") or context.get("input")
        if not input_data:
            raise ValueError("Chemist context missing 'smiles' or 'input'")

        result = self.process_molecule(input_data, context=user_context)

        return {
            "status": "success",
            "content": f"Processed molecule: {result.canonical_smiles}",
            "artifacts": [{"type": "molecule", "data": result.model_dump(mode="json")}],
        }

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

        Used by downstream agents (Cortex) to verify generated molecules.

        Args:
            structure_string (str): SMILES or SELFIES string.
            context (UserContext): The user identity context.

        Returns:
            MolecularEntity: A model with computed properties if valid,
            or status=CORRUPT if invalid.
        """
        return Chemist.process_molecule(structure_string, context=context)

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

        Args:
            input_data (Union[str, Path]): SMILES/InChI/SELFIES string or path to a molecule file.
            context (UserContext): The user identity context.

        Returns:
            MolecularEntity: A model populated with computed properties.
            If processing fails, returns a MolecularEntity with status=CORRUPT.

        Raises:
            ValueError: If molecule sanitization fails or processing encounters
            an unrecoverable error.
        """
        logger.info(f"Processing molecule input: {input_data}")
        input_str_for_hash = str(input_data)
        logger.info(
            "Analyzing molecular structure",
            user_id=context.user_id.get_secret_value(),
            smiles_hash=hash(input_str_for_hash),
        )
        mol = None
        source_smiles = ""

        # 1. Load Molecule
        try:
            if isinstance(input_data, Path) or (isinstance(input_data, str) and Path(input_data).exists()):
                # datamol doesn't have read_mol, use read_sdf for .sdf or specific readers, or fallback to rdkit
                # Since we want to handle generic files, we can check extension
                file_path = str(input_data)
                ext = Path(file_path).suffix.lower()
                if ext == ".mol":
                    mol = dm.read_molblock(Path(file_path).read_text())
                elif ext == ".sdf":
                    # read_sdf returns a list, take first
                    mols = dm.read_sdf(file_path)
                    if mols:
                        mol = mols[0]
                else:
                    # Fallback to smart reading if available or treat as error for now
                    raise ValueError(f"Unsupported file extension: {ext}")

                if mol:
                    source_smiles = dm.to_smiles(mol)
            else:
                input_str = str(input_data).strip()
                # Detection strategy:
                # 1. Check for InChI prefix
                if input_str.startswith("InChI="):
                    logger.info("Detected InChI input")
                    # datamol.to_mol does not support InChI strings directly, use RDKit
                    mol = Chem.MolFromInchi(input_str)
                    if mol:
                        source_smiles = dm.to_smiles(mol)
                else:
                    # 2. Check if it's SELFIES
                    # There is no perfect regex for SELFIES, but they typically contain [...] blocks.
                    # Or we can try to decode it as SELFIES first if it looks like one, or just try to convert.
                    # A robust way is to try decoding if it contains typical SELFIES chars like '[' and ']'.
                    # However, SMILES also have brackets.
                    # Let's try to interpret as SMILES first (default behavior).
                    # If that fails, try SELFIES.
                    # OR if we want to be explicit, maybe the caller should specify, but requirements say "Raw string".

                    # Let's try dm.to_mol (SMILES) first.
                    source_smiles = input_str
                    mol = dm.to_mol(source_smiles)

                    if mol is None:
                        # Try interpreting as SELFIES
                        try:
                            logger.info("SMILES parsing failed, attempting SELFIES decoding")
                            decoded_smiles = sf.decoder(input_str)
                            if decoded_smiles:
                                mol = dm.to_mol(decoded_smiles)
                                if mol:
                                    source_smiles = (
                                        decoded_smiles  # We treat the decoded SMILES as source for processing
                                    )
                        except Exception:
                            # Not a valid SELFIES either
                            pass

            if mol is None:
                logger.error(f"Failed to parse molecule from input: {input_data}")
                raise ValueError(f"Invalid molecule input: {input_data}")

            # 2. Normalization (Sanitization & Canonicalization)

            # sanitizes and removes salts, solvents, neutralizes charges
            # PRD says "Removes salts, solvents, and neutralizes charges to find the 'Parent Structure.'"

            # Step 2a: Remove salts/solvents
            # dm.remove_salts_solvents uses a list of common salts.
            # dont_remove_everything=True prevents removing the drug itself if it happens to be in the
            # salt list (e.g. Benzoic Acid).
            mol = dm.remove_salts_solvents(mol, dont_remove_everything=True)

            # Step 2b: Keep Largest Fragment
            # Ensure we have a single parent structure if multiple fragments remain (e.g. non-salt mixture).
            mol = dm.keep_largest_fragment(mol)

            # Step 2c: Standardize (includes neutralization/uncharging)
            # using standardize_mol with uncharge=True seems more robust than sanitize_mol(charge_neutral=True)
            # for converting [O-] to O (adding H).
            mol = dm.standardize_mol(mol, uncharge=True)

            # Step 2d: Sanitize to be sure (standardize might return unsanitized mol in some edge cases?)
            # But standardize usually returns a sanitized mol.
            # Let's just do a final sanity check or lightweight sanitize if needed.
            # dm.sanitize_mol is still good for fixing aromaticity if standardize didn't catch it.
            mol = dm.sanitize_mol(mol, sanifix=True, verbose=False)  # charge_neutral=False as standardize did it

            if mol is None:
                raise ValueError("Molecule sanitization resulted in None")

            canonical_smiles = dm.to_smiles(mol, canonical=True)

            # 3. Transmutation (SMILES -> SELFIES)
            try:
                selfies_string = sf.encoder(canonical_smiles)
            except Exception as e:
                logger.error(f"SELFIES conversion failed for: {canonical_smiles}. Error: {e}")
                # Fallback or raise? PRD says "Transmutation ... Why SOTA: SELFIES are 100% valid."
                # If we can't generate it, it's an issue.
                raise ValueError("SELFIES conversion failed") from e

            # 4. Computed Properties
            # Molecular Weight
            mw = dm.descriptors.mw(mol)
            # LogP
            logp = dm.descriptors.clogp(mol)
            # TPSA
            tpsa = dm.descriptors.tpsa(mol)
            # Rotatable Bonds
            rotatable_bonds = dm.descriptors.n_rotatable_bonds(mol)

            # Lipinski Violations
            # Rules: MW <= 500, LogP <= 5, HBD <= 5, HBA <= 10
            hbd = dm.descriptors.n_lipinski_hbd(mol)
            hba = dm.descriptors.n_lipinski_hba(mol)

            violations = 0
            if mw > 500:
                violations += 1
            if logp > 5:
                violations += 1
            if hbd > 5:
                violations += 1
            if hba > 10:
                violations += 1

            # 5. Fingerprint (Morgan/ECFP)
            # Using radius 2 (ECFP4) and 2048 bits as standard defaults
            fp = dm.to_fp(mol, as_array=True, fp_type="ecfp", radius=2, fpSize=2048)
            fingerprint_vector: List[int] = fp.tolist()

            # 6. Universal Hash (InChI Key)
            inchi_key = dm.to_inchikey(mol)

            return MolecularEntity(
                id=uuid4(),
                source_smiles=source_smiles,
                status=ProcessingStatus.VALID,
                canonical_smiles=canonical_smiles,
                selfies_string=selfies_string,
                inchi_key=inchi_key,
                molecular_weight=float(mw),
                logp=float(logp),
                tpsa=float(tpsa),
                num_rotatable_bonds=int(rotatable_bonds),
                lipinski_violations=violations,
                fingerprint_vector=fingerprint_vector,
            )

        except Exception as e:
            logger.error(f"Processing failed for input: {input_data}. Error: {e}")
            return MolecularEntity(
                id=uuid4(),
                source_smiles=source_smiles if source_smiles else str(input_data),
                status=ProcessingStatus.CORRUPT,
            )
