src/ui/handlers.py

"""
UI Handlers Module

This module contains event handlers and business logic
for the drug discovery application UI components.
"""

from ..molecules.analysis import (
    analyze_molecule_image_only,
    validate_smiles_realtime,
    get_molecule_properties_for_hover,
)
from ..molecules.generated_variations import generate_variations_from_partial_smiles
from ..ai.services import respond, handle_structure_chat, parse_ai_structures


class VariationHandlers:
    """Handles variation-related functionality."""
    
    def __init__(self):
        self.current_variations = []      # list of dicts: {smiles, image, style}
        self.current_page = 0
        self.variations_per_page = 12

    # -------------------------------------------------------------
    # Generate variations for the UI grid
    # -------------------------------------------------------------
    def generate_variations_for_display(self, input_smiles, num_variations=12):
        """
        Generate variations using Gen_PartialSMILES2.py, convert to RDKit images,
        and return formatted gallery items while printing debugging info.
        """

        print("\n==============================")
        print("🚀 generate_variations_for_display CALLED")
        print("==============================")
        print(f"User input SMILES: {input_smiles}")
        print(f"Requested # variations: {num_variations}")

        # ------------------------------------------------------------
        # Create big image for the user input molecule
        # ------------------------------------------------------------
        from rdkit import Chem
        from rdkit.Chem import Draw

        mol_input = Chem.MolFromSmiles(input_smiles)
        if mol_input is None:
            print("❌ ERROR: Invalid SMILES input. Cannot generate big preview.")
            return [], input_smiles, "Invalid SMILES"

        big_image = Draw.MolToImage(mol_input, size=(400, 300))

        print("✔ Input molecule rendered successfully.")

        # ------------------------------------------------------------
        # Run the generation helper
        # ------------------------------------------------------------
        variations = generate_variations_from_partial_smiles(
            input_smiles,
            n_to_gen=num_variations
        )

        print(f"Generator returned {len(variations)} variations.")

        # No variations returned ≠ UI failure — show empty grid
        if not variations:
            print(" No variations generated — returning empty gallery.")
            return [], input_smiles, ""

        self.current_variations = variations

        # ------------------------------------------------------------
        # Build gallery items for Gradio
        # ------------------------------------------------------------
        gallery_items = []
        print("\n Building gallery_items...\n")

        for i, v in enumerate(variations):
            img = v["image"]
            smi = v["smiles"]

            print(f"  #{i+1}:")
            print(f"     SMILES: {smi}")
            print(f"     Image type: {type(img)}")
            try:
                print(f"Image size: {img.size}")
            except:
                print("Could not read image size")

            gallery_items.append((img, smi))

        print("\n==============================")
        print("Returning data to UI")
        print(f"Total gallery items: {len(gallery_items)}")
        print("==============================\n")

        # ------------------------------------------------------------
        # Return to Gradio exactly as before
        # ------------------------------------------------------------
        return gallery_items, input_smiles, ""

    # -------------------------------------------------------------
    # Handle selection from the variations grid
    # -------------------------------------------------------------
    def select_variation(self, evt):
        """Handle selection of a variation from the grid."""
        try:
            print("=== SELECT_VARIATION CALLED ===")
            print(f"Event: {evt}, type: {type(evt)}")
            print(f"Current variations count: {len(self.current_variations)}")

            if not self.current_variations:
                return None, "", "", ""

            # Gradio events might carry an index, or evt might be None/int
            if evt is None:
                index = 0
            elif hasattr(evt, "index"):
                index = evt.index
            elif isinstance(evt, (int, float)):
                index = int(evt)
            else:
                index = 0

            # Clamp index
            if index < 0 or index >= len(self.current_variations):
                index = 0

            selected_var = self.current_variations[index]
            smiles = selected_var["smiles"]
            print(f"Selected SMILES: {smiles}")

            properties_text = get_molecule_properties_for_hover(smiles)

            # Outputs:
            #   main_structure_display,
            #   selected_smiles_display,
            #   selected_style_display,
            #   properties_display
            return selected_var["image"], smiles, selected_var["style"], properties_text

        except Exception as e:
            print(f"Error in select_variation: {e}")
            import traceback
            traceback.print_exc()
            return None, "", "", ""

    # -------------------------------------------------------------
    # Clear variations
    # -------------------------------------------------------------
    def clear_variations(self):
        """Clear all variations and reset display."""
        self.current_variations = []
        self.current_page = 0
        return [], "", ""

    # -------------------------------------------------------------
    # Pagination (prev/next buttons)
    # -------------------------------------------------------------
    def navigate_variations(self, direction):
        """Navigate through variations pages."""
        if not self.current_variations:
            return [], "Page 1 of 1", None, "", ""

        total = len(self.current_variations)
        pages = (total + self.variations_per_page - 1) // self.variations_per_page

        if direction == "next":
            self.current_page = min(self.current_page + 1, pages - 1)
        elif direction == "prev":
            self.current_page = max(self.current_page - 1, 0)

        start = self.current_page * self.variations_per_page
        end = min(start + self.variations_per_page, total)

        page_variations = self.current_variations[start:end]
        gallery_items = [(v["image"], v["smiles"]) for v in page_variations]

        page_info = f"Page {self.current_page + 1} of {pages}"

        first = page_variations[0] if page_variations else None

        return (
            gallery_items,
            page_info,
            first["image"] if first else None,
            first["smiles"] if first else "",
            first["style"] if first else "",
        )

    # -------------------------------------------------------------
    # Update how many variations per page
    # -------------------------------------------------------------
    def update_variation_count(self, count):
        """Update the number of variations to generate and show per page."""
        self.variations_per_page = count
        return count

    # -------------------------------------------------------------
    # Analyze molecule + tooltip text (used on load / selection)
    # -------------------------------------------------------------
    def analyze_molecule_with_tooltip(self, smiles):
        """Analyze molecule and return RDKit image and tooltip text."""
        molecule_img = analyze_molecule_image_only(smiles)
        tooltip_text = get_molecule_properties_for_hover(smiles)
        return molecule_img, tooltip_text


class BookmarkHandlers:
    """Handles bookmark functionality."""
    
    def __init__(self):
        self.bookmarked_molecules = []
    
    def bookmark_molecule(self, smiles, molecule_name=""):
        """Add a molecule to the bookmarked collection."""
        from rdkit import Chem
        from rdkit.Chem import Draw
        
        mol = Chem.MolFromSmiles(smiles)
        if not mol:
            return "❌ Invalid SMILES string - cannot bookmark"
        
        if smiles in [bm['smiles'] for bm in self.bookmarked_molecules]:
            return "⚠️ Molecule already bookmarked"
        
        if not molecule_name:
            molecule_name = f"Bookmarked_{len(self.bookmarked_molecules) + 1}"
        
        self.bookmarked_molecules.append({
            'smiles': smiles,
            'name': molecule_name,
            'timestamp': len(self.bookmarked_molecules) + 1
        })
        
        return f"✅ Bookmarked: {molecule_name}"
    
    def get_bookmarked_molecules(self):
        return self.bookmarked_molecules
    
    def remove_bookmark(self, smiles):
        self.bookmarked_molecules = [bm for bm in self.bookmarked_molecules if bm['smiles'] != smiles]
        return "🗑️ Removed from bookmarks"
    
    def bookmark_current_molecule(self, smiles, name):
        """Bookmark current molecule and update gallery."""
        from rdkit import Chem
        from rdkit.Chem import Draw
        
        result = self.bookmark_molecule(smiles, name)
        bookmarked_mols = self.get_bookmarked_molecules()
        gallery_items = []
        for mol in bookmarked_mols:
            mol_obj = Chem.MolFromSmiles(mol['smiles'])
            if mol_obj:
                img = Draw.MolToImage(mol_obj, size=(150, 150), kekulize=True)
                gallery_items.append((img, f"{mol['name']}: {mol['smiles']}"))
        return result, gallery_items


class AIHandler:
    """Handles AI chat functionality with both general questions and structure generation."""
    
    def __init__(self):
        self.chat_history = []
    
    def handle_ai_chat(self, message, history, selected_smiles, hf_token, temperature):
        """Handle AI chat with both general questions and structure generation."""
        if not message.strip() or not hf_token.strip():
            return history, []
        
        history.append({"role": "user", "content": message})
        
        structure_keywords = ['generate', 'create', 'modify', 'derivative', 'variant', 'structure']
        is_structure_request = any(keyword in message.lower() for keyword in structure_keywords)
        
        if is_structure_request and selected_smiles:
            ai_response = ""
            for chunk in respond(
                message,
                history[:-1],
                "You are an expert medicinal chemist. Generate new chemical structures based on user requests.",
                512,
                temperature,
                0.9,
                hf_token
            ):
                ai_response = chunk
            
            history.append({"role": "assistant", "content": ai_response})
            structures = parse_ai_structures(ai_response, selected_smiles)
            return history, structures
        else:
            ai_response = ""
            for chunk in respond(
                message,
                history[:-1],
                "You are an expert medicinal chemist and drug discovery specialist. Help with molecular analysis, drug design, and medicinal chemistry questions.",
                512,
                temperature,
                0.9,
                hf_token
            ):
                ai_response = chunk
            
            history.append({"role": "assistant", "content": ai_response})
            return history, []