from langchain_groq import ChatGroq
import os
from .models import Sections

# Global LLM instance
llm = None

def initialize_llm(groq_api_key=None):
    """
    Initialize the LLM with optional API key
    
    Args:
        groq_api_key (str, optional): Groq API key to use
    """
    global llm
    
    # Set API key if provided
    if groq_api_key:
        os.environ["GROQ_API_KEY"] = groq_api_key
        
    # Initialize LLM
    llm = ChatGroq(model="qwen-2.5-32b")
    
    return llm

def get_llm():
    """
    Get the global LLM instance, initializing if needed
    
    Returns:
        ChatGroq: The LLM instance
    """
    global llm
    
    if llm is None:
        llm = initialize_llm()
        
    return llm

def get_planner():
    """
    Get the planner with structured output capabilities
    
    Returns:
        Structured LLM: LLM with structured output for planning
    """
    # Set up structured output model for planner
    return get_llm().with_structured_output(Sections)