from .utils import format_prompt_for_hint, format_prompt_for_followup
from .model_handler import generate_hint, generate_follow_up, generate_feedback # Add generate_feedback
import re # Import regex module

def clean_task_description(description: str) -> str:
    """Removes common example, constraint, input/output format sections by finding the first occurrence."""
    lines = description.splitlines()
    cleaned_lines = []
    skip_keywords = ["Example:", "Input:", "Output:", "Constraints:", "Input Format:", "Output Format:", "Explanation:"]
    first_skip_index = -1

    # Find the index of the first line starting with a skip keyword
    for i, line in enumerate(lines):
        line_stripped_lower = line.strip().lower()
        if any(line_stripped_lower.startswith(keyword.lower()) for keyword in skip_keywords):
            first_skip_index = i
            break # Stop at the first occurrence

    # If a skip keyword was found, take lines before it
    if first_skip_index != -1:
        cleaned_lines = lines[:first_skip_index]
    # Otherwise, keep all lines (no skip keywords found)
    else:
        cleaned_lines = lines

    # Join back the relevant lines and strip whitespace
    return "\n".join(cleaned_lines).strip()


def get_hint(code, task_description, mode='concise'): # Add mode
    cleaned_description = clean_task_description(task_description)
    # Pass mode to generate_hint
    return generate_hint(code, cleaned_description, mode)

def get_feedback(code, task_description, mode='concise'): # New function
    cleaned_description = clean_task_description(task_description)
    # Call generate_feedback with mode
    return generate_feedback(code, cleaned_description, mode)

def get_follow_up_question(code, task_description):
    cleaned_description = clean_task_description(task_description)
    # Pass code to generate_follow_up (already updated in model_handler)
    return generate_follow_up(cleaned_description, code) # Pass both args