"""Organic structure diagram agent using chemfig.

Generates molecular structures for organic chemistry using chemfig package.
"""

from typing import Optional

from vbagent.agents.base import (
    create_agent,
    create_image_message,
    run_agent_sync,
)
from vbagent.prompts.diagram.chemistry.organic_structure import (
    SYSTEM_PROMPT,
    USER_TEMPLATE,
    USER_TEMPLATE_MCQ_OPTIONS,
    USER_TEMPLATE_FROM_PROBLEM,
)
from vbagent.references.store import ReferenceStore
from vbagent.utils.latex import clean_latex_output


def create_organic_structure_agent(use_context: bool = True) -> "Agent":
    """Create an organic structure generation agent.
    
    Args:
        use_context: Whether to include reference context
        
    Returns:
        Configured Agent instance
    """
    prompt = SYSTEM_PROMPT
    
    # Add reference context if enabled
    if use_context:
        context = get_organic_structure_context_for_classification()
        if context:
            prompt = prompt + "\n\n" + context
    
    return create_agent(
        name="OrganicStructure",
        instructions=prompt,
        agent_type="tikz",
    )


def generate_organic_structure(
    image_path: Optional[str] = None,
    description: Optional[str] = None,
    use_context: bool = True,
    show_spinner: bool = True,
    mcq_options: bool = False,
) -> str:
    """Generate chemfig code for an organic structure.
    
    Args:
        image_path: Path to structure image (optional)
        description: Text description of structure (optional)
        use_context: Whether to include reference context
        show_spinner: Whether to show animated spinner
        mcq_options: Whether to generate all 4 MCQ options (A, B, C, D)
        
    Returns:
        chemfig code as string (or \\def\\OptionA{...} blocks if mcq_options=True)
        
    Raises:
        ValueError: If neither image_path nor description provided
    """
    if not image_path and not description:
        raise ValueError("Either image_path or description must be provided")
    
    agent = create_organic_structure_agent(use_context)
    
    if image_path:
        # Choose template based on mcq_options flag
        template = USER_TEMPLATE_MCQ_OPTIONS if mcq_options else USER_TEMPLATE
        message = create_image_message(image_path, template)
    else:
        # For text-only description, create a simple text message
        message = [{"role": "user", "content": USER_TEMPLATE_FROM_PROBLEM.format(problem=description)}]
    
    raw_output = run_agent_sync(agent, message, show_spinner=show_spinner)
    return clean_latex_output(raw_output)


def validate_organic_structure_output(chemfig_code: str) -> tuple[bool, str]:
    """Validate chemfig code for organic structures.
    
    Args:
        chemfig_code: The chemfig code to validate
        
    Returns:
        Tuple of (is_valid, error_message)
    """
    if not chemfig_code or not chemfig_code.strip():
        return False, "Empty chemfig code"
    
    # Check for chemfig command
    if "\\chemfig{" not in chemfig_code:
        return False, "Missing \\chemfig{} command"
    
    # Check for balanced braces
    open_count = chemfig_code.count("{")
    close_count = chemfig_code.count("}")
    if open_count != close_count:
        return False, f"Unbalanced braces: {open_count} open, {close_count} close"
    
    # Check for balanced parentheses (branches)
    open_paren = chemfig_code.count("(")
    close_paren = chemfig_code.count(")")
    if open_paren != close_paren:
        return False, f"Unbalanced parentheses: {open_paren} open, {close_paren} close"
    
    # Check for TikZ commands (should not be present)
    tikz_commands = ["\\begin{tikzpicture}", "\\draw", "\\node", "\\coordinate"]
    for cmd in tikz_commands:
        if cmd in chemfig_code:
            return False, f"Should use chemfig, not TikZ commands like {cmd}"
    
    return True, ""


def get_organic_structure_context_for_classification() -> str:
    """Get reference context for organic structure generation.
    
    Returns:
        Context string with relevant examples
    """
    store = ReferenceStore()
    
    # Search for organic chemistry structure examples
    results = store.search(
        query="organic structure chemfig benzene alkane functional group",
        file_types=["sty", "tex", "pdf"],
        max_results=3
    )
    
    if not results:
        return ""
    
    context = "\n\n## Reference Examples\n\n"
    context += "Here are some similar organic structures from our reference library:\n\n"
    
    for i, result in enumerate(results, 1):
        context += f"### Example {i}\n"
        context += f"```latex\n{result.content}\n```\n\n"
    
    return context
