src/code_generation.py

import os, re, time
from openai import OpenAI
import google.generativeai as genai
from dotenv import load_dotenv
import logging
from typing import Dict, Optional
from src.model_config import GROQ_CONTENT_MODELS, get_client_for_task, get_model_for_task, get_service_for_task

load_dotenv()

logger = logging.getLogger(__name__)

# API Keys
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

# Groq models to try
GROQ_MODELS = GROQ_CONTENT_MODELS

# Initialize clients
openai_client = None
groq_client = None
openrouter_client = None

if OPENAI_API_KEY:
    openai_client = OpenAI(api_key=OPENAI_API_KEY)
    logger.info("OpenAI client initialized for code generation")

if GROQ_API_KEY:
    groq_client = OpenAI(
        api_key=GROQ_API_KEY,
        base_url="https://api.groq.com/openai/v1"
    )
    logger.info("Groq client initialized for code generation (llama-3.3-70b-versatile)")

# Set primary client (priority: OpenAI > Groq)
client = openai_client if openai_client else groq_client


def _call_llm_with_fallback(messages, model=None, temperature=0.5, max_tokens=500, task_name="content_generation"):
    """Call LLM with fallback from task-specific client -> OpenAI -> OpenRouter."""
    # Determine task-specific client and model if not provided
    if model is None:
        try:
            model = get_model_for_task(task_name)
            task_client = get_client_for_task(task_name)
            service = get_service_for_task(task_name)
            logger.info(f"Using {service} ({model}) for task: {task_name}")
        except Exception:
            task_client = None
            service = None
    else:
        task_client = None

        # If a task-specific client (e.g., Anthropic) is available, try it first
        if task_client:
            try:
                # Check if it's Anthropic client
                if service == "Anthropic":
                    # Anthropic uses messages.create with system as separate parameter
                    system_msg = next((m['content'] for m in messages if m['role'] == 'system'), None)
                    user_messages = [m for m in messages if m['role'] != 'system']
                    
                    if system_msg:
                        anthro_response = task_client.messages.create(
                            model=model,
                            max_tokens=max_tokens,
                            temperature=temperature,
                            system=system_msg,
                            messages=user_messages
                        )
                    else:
                        anthro_response = task_client.messages.create(
                            model=model,
                            max_tokens=max_tokens,
                            temperature=temperature,
                            messages=user_messages
                        )
                    # Convert to OpenAI format
                    class MockResponse:
                        def __init__(self, content):
                            self.choices = [type('obj', (object,), {'message': type('obj', (object,), {'content': content})})()]
                    response = MockResponse(anthro_response.content[0].text)
                else:
                    response = task_client.chat.completions.create(
                        model=model,
                        messages=messages,
                        temperature=temperature,
                        max_tokens=max_tokens,
                    )
                return response
            except Exception as e:
                logger.warning(f"Task-specific client ({service}) failed: {e}. Trying OpenAI/OpenRouter...")    # Try OpenAI client if configured
    if openai_client:
        try:
            response = openai_client.chat.completions.create(
                model=model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
            )
            return response
        except Exception as e:
            logger.warning(f"OpenAI failed: {e}. Falling back to Gemini...")

    # Fallback to Gemini (if OpenAI failed)
    try:
        logger.warning("OpenAI failed. Falling back to Gemini...")
        prompt = "\n".join([m["content"] for m in messages])
        gemini_model = genai.GenerativeModel("gemini-2.0-flash")
        gemini_response = gemini_model.generate_content(prompt)

        # Convert Gemini response to OpenAI-like format
        class MockResponse:
            def __init__(self, content):
                self.choices = [
                    type(
                        "obj",
                        (object,),
                        {"message": type("obj", (object,), {"content": content})},
                     )()
                ]
        return MockResponse(gemini_response.text)
    except Exception as e:
        logger.error(f"Gemini fallback failed: {e}")
        raise RuntimeError("All LLM providers failed (OpenAI and Gemini unavailable)")

def generate_code_example(
    topic: str, programming_language: str, context: str
) -> Optional[Dict[str, str]]:
    """
    Generates a code example and its explanation based on the presentation context.

    This function performs two main actions:
    1. Generates a concise and relevant code snippet for the given topic.
    2. Generates a simple, clear explanation for that code.

    Args:
        topic (str): The main topic of the presentation.
        programming_language (str): The programming language to be used.
        context (str): The broader text from the presentation for context.

    Returns:
        A dictionary containing the 'title', 'language', 'code', and 'explanation',
        or None if the generation fails.
    """
    if not client:
        logger.error("No LLM client is initialized. Cannot generate code.")
        return None

    try:
        code_generation_prompt = f"""
        Based on the following presentation context about '{topic}', generate a single, concise, and runnable code example in {programming_language}.

        Context:
        ---
        {context}
        ---
        Requirements:
        1.  Code Only: Your response must contain ONLY the code, wrapped in a single markdown block (e.g., ```{programming_language}\n...\n```).
        2.  Concise: The code should be between 4 and 10 lines.
        3.  Relevant: The code must be a direct and clear example of the main topic.
        4.  No Explanation: Do not add any explanation or introductory text in this response.
        5.  No thinking tags or extra text.
        """
        code_completion = _call_llm_with_fallback(
            messages=[
                {"role": "system", "content": "You are a senior software engineer who writes clear, exemplary code snippets. Output ONLY code in markdown blocks."},
                {"role": "user", "content": code_generation_prompt},
            ],
            temperature=0.3,
            max_tokens=500,
        )
        raw_code = code_completion.choices[0].message.content.strip()
        
        # Clean up thinking tags if present
        raw_code = re.sub(r'<think>.*?</think>', '', raw_code, flags=re.DOTALL).strip()
        
        code_match = re.search(f"```{programming_language}\n(.*?)\n```", raw_code, re.DOTALL)
        if not code_match:
            code_match = re.search(r"```\w*\n(.*?)\n```", raw_code, re.DOTALL)
        if not code_match:
            code_match = re.search(r"```\n(.*?)\n```", raw_code, re.DOTALL)
        if not code_match:
            logger.error("Could not parse the generated code snippet from the LLM response.")
            return None
        code_snippet = code_match.group(1).strip()

        explanation_prompt = f"""
        Generate a simple, educational explanation for the following {programming_language} code snippet.
        Code Snippet:
        ```
        {code_snippet}
        ```
        Requirements:
        1.  Title: Start with a concise title for the slide, prefixed with "Title: ".
        2.  Explanation: Provide a brief, clear explanation of what the code does. Explain it as you would to a beginner.
        3.  Format: The explanation should be a short paragraph or a few bullet points.
        4.  Concise: Keep the entire response under 100 words.
        5.  No thinking tags or extra text.
        """
        explanation_completion = _call_llm_with_fallback(
            messages=[
                {"role": "system", "content": "You are a friendly and clear technical educator. Output ONLY the explanation."},
                {"role": "user", "content": explanation_prompt},
            ],
            temperature=0.7,
            max_tokens=300,
        )
        explanation_response = explanation_completion.choices[0].message.content.strip()
        
        # Clean up thinking tags if present
        explanation_response = re.sub(r'<think>.*?</think>', '', explanation_response, flags=re.DOTALL).strip()

        title_match = re.search(r"Title:\s*(.*)", explanation_response)
        title = title_match.group(1).strip() if title_match else "Code Breakdown"
        explanation = explanation_response.replace(f"Title: {title}", "").strip()
        logger.info(f"Successfully generated code example and explanation for '{topic}'.")
        return {
            "title": title,
            "language": programming_language,
            "code": code_snippet,
            "explanation": explanation,
        }
    except Exception as e:
        logger.error(f"An error occurred during code generation: {e}", exc_info=True)
        return None