"""
AI Services Module

This module handles AI-powered features including
chat responses and structure generation.
"""

from huggingface_hub import InferenceClient
from rdkit import Chem
from rdkit.Chem import Draw


def respond(
    message,
    history: list[dict[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    hf_token: str,
):
    """
    Enhanced drug discovery chatbot with molecular property integration.
    """
    client = InferenceClient(token=hf_token, model="openai/gpt-oss-20b")

    # Enhanced system message for drug discovery
    enhanced_system_message = f"""{system_message}

You are an expert medicinal chemist and drug discovery specialist. You help researchers understand:
- Molecular properties and drug-likeness
- Structure-activity relationships (SAR)
- ADMET properties (Absorption, Distribution, Metabolism, Excretion, Toxicity)
- Drug design principles and optimization strategies
- Chemical synthesis and medicinal chemistry

When discussing molecules, consider their molecular weight, LogP, hydrogen bonding, and other key properties. Reference the molecular gallery below for visual context."""

    messages = [{"role": "system", "content": enhanced_system_message}]
    messages.extend(history)
    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        choices = message.choices
        token = ""
        if len(choices) and choices[0].delta.content:
            token = choices[0].delta.content

        response += token
        yield response


def generate_ai_structures(message, history, selected_smiles, hf_token: str):
    """Generate new structures using AI based on user request and selected molecule."""
    if not message.strip():
        return history, []
    
    client = InferenceClient(token=hf_token, model="openai/gpt-oss-20b")
    
    # Enhanced system message for structure generation
    system_message = f"""You are an expert medicinal chemist and drug discovery specialist. You help generate new chemical structures based on user requests and existing molecules.

Current selected molecule SMILES: {selected_smiles}

Your task is to:
1. Understand the user's request for structure modifications
2. Generate 3-6 new SMILES strings that meet their requirements
3. Provide brief explanations for each generated structure
4. Consider drug-likeness, ADMET properties, and medicinal chemistry principles

Format your response as:
STRUCTURE 1: [SMILES] - [Brief explanation]
STRUCTURE 2: [SMILES] - [Brief explanation]
etc.

Focus on practical, synthesizable structures that address the user's specific request."""

    messages = [{"role": "system", "content": system_message}]
    messages.extend(history)
    messages.append({"role": "user", "content": message})

    response = ""
    for message_chunk in client.chat_completion(
        messages,
        max_tokens=512,
        stream=True,
        temperature=0.7,
        top_p=0.9,
    ):
        choices = message_chunk.choices
        token = ""
        if len(choices) and choices[0].delta.content:
            token = choices[0].delta.content
        response += token
        yield response


def parse_ai_structures(ai_response, selected_smiles):
    """Parse AI response to extract SMILES strings and generate images."""
    structures = []
    lines = ai_response.split('\n')
    
    for line in lines:
        if 'STRUCTURE' in line.upper() and ':' in line:
            try:
                # Extract SMILES from lines like "STRUCTURE 1: CCO - ethanol"
                parts = line.split(':', 1)
                if len(parts) > 1:
                    smiles_part = parts[1].split('-')[0].strip()
                    # Clean up the SMILES string
                    smiles = smiles_part.strip('[]()').strip()
                    
                    # Validate and generate image
                    mol = Chem.MolFromSmiles(smiles)
                    if mol:
                        img = Draw.MolToImage(mol, size=(150, 150), kekulize=True)
                        structures.append((img, f"Generated: {smiles}"))
            except:
                continue
    
    # If no structures were parsed, generate some variations of the selected molecule
    if not structures and selected_smiles:
        try:
            mol = Chem.MolFromSmiles(selected_smiles)
            if mol:
                # Generate a few variations with different sizes
                for size in [(120, 120), (150, 150), (180, 180)]:
                    img = Draw.MolToImage(mol, size=size, kekulize=True)
                    structures.append((img, f"Variation: {selected_smiles}"))
        except:
            pass
    
    return structures


def handle_structure_chat(message, history, selected_smiles, hf_token: str):
    """Handle the structure generation chat."""
    if not message.strip():
        return history, []
    
    # Add user message to history
    history.append({"role": "user", "content": message})
    
    # Generate AI response
    ai_response = ""
    for chunk in generate_ai_structures(message, history[:-1], selected_smiles, hf_token):
        ai_response = chunk
    
    # Add AI response to history
    history.append({"role": "assistant", "content": ai_response})
    
    # Parse and generate structure images
    structures = parse_ai_structures(ai_response, selected_smiles)
    
    return history, structures