code_router.py

import re

from schemas import CodeTaskType


GENERATE_PATTERNS = [
    r"\bcreate\b",
    r"\bgenerate\b",
    r"\bbuild\b",
    r"\bwrite\b",
    r"\bmake\b",
    r"\bdevelop\b",
    r"\bimplement\b",
    r"\bcreate from scratch\b",
    r"\bstart a new\b",
]


FIX_PATTERNS = [
    r"\bfix\b",
    r"\bsolve\b",
    r"\bcorrect\b",
    r"\brepair\b",
    r"\bnot working\b",
    r"\bissue\b",
    r"\bbug\b",
    r"\berror\b",
    r"\bexception\b",
    r"\bcrash\b",
    r"\bfailed\b",
    r"\bproblem\b",
    r"\bdoesn't work\b",
    r"\bdoes not work\b",
    r"\bthrowing\b",
    r"\btraceback\b",
]


EXPLAIN_PATTERNS = [
    r"\bexplain\b",
    r"\bwhat does this do\b",
    r"\bhow does this work\b",
    r"\bdescribe\b",
    r"\bmeaning\b",
    r"\bunderstand\b",
    r"\bwalk me through\b",
    r"\bcan you explain\b",
]


REFACTOR_PATTERNS = [
    r"\brefactor\b",
    r"\bclean\b",
    r"\bimprove\b",
    r"\boptimize\b",
    r"\bbetter\b",
    r"\bmake this cleaner\b",
    r"\bmake this better\b",
    r"\bclean up\b",
]


REVIEW_PATTERNS = [
    r"\breview\b",
    r"\bcheck this code\b",
    r"\baudit\b",
    r"\binspect\b",
    r"\bcode review\b",
    r"\bany issues\b",
    r"\bis this okay\b",
    r"\bfind problems\b",
]


MODIFICATION_PATTERNS = [
    r"\badd\b",
    r"\bupdate\b",
    r"\bmodify\b",
    r"\bchange\b",
    r"\brename\b",
    r"\bconvert\b",
    r"\bextend\b",
    r"\binclude\b",
    r"\buse\b",
    r"\breplace\b",
    r"\bremove\b",
    r"\bdelete\b",
    r"\bdocstring\b",
    r"\btype hints?\b",
    r"\bvalidation\b",
    r"\berror handling\b",
    r"\blogging\b",
    r"\bcomments\b",
    r"\bexception handling\b",
    r"\bmake it\b",
    r"\bturn this into\b",
    r"\badd support for\b",
]


DIRECT_EXPLAIN_ONLY_PATTERNS = [
    r"\bexplain\b",
    r"\bwhat does this do\b",
    r"\bhow does this work\b",
    r"\bdescribe this code\b",
    r"\bhelp me understand\b",
    r"\bwalk me through\b",
]


def normalize_text(text: str) -> str:
    text = text or ""
    text = text.strip().lower()
    text = re.sub(r"\s+", " ", text)
    return text


def contains_pattern(text: str, patterns: list[str]) -> bool:
    return any(re.search(pattern, text) for pattern in patterns)


def is_modification_request(text: str) -> bool:
    return contains_pattern(text, MODIFICATION_PATTERNS)


def is_direct_explanation_request(text: str) -> bool:
    return contains_pattern(text, DIRECT_EXPLAIN_ONLY_PATTERNS)


def detect_task_type(
    message: str,
    code: str | None = None,
    error_message: str | None = None,
    mode_hint: CodeTaskType | None = None,
) -> CodeTaskType:
    if mode_hint and mode_hint != CodeTaskType.UNKNOWN:
        return mode_hint

    normalized_message = normalize_text(message)
    has_code = bool(code and code.strip())
    has_error = bool(error_message and error_message.strip())

    if has_error:
        return CodeTaskType.FIX

    if contains_pattern(normalized_message, FIX_PATTERNS):
        return CodeTaskType.FIX

    if contains_pattern(normalized_message, REVIEW_PATTERNS):
        return CodeTaskType.REVIEW

    if contains_pattern(normalized_message, REFACTOR_PATTERNS):
        return CodeTaskType.REFACTOR

    if is_direct_explanation_request(normalized_message) and not is_modification_request(normalized_message):
        return CodeTaskType.EXPLAIN

    if has_code and is_modification_request(normalized_message):
        return CodeTaskType.REFACTOR

    if contains_pattern(normalized_message, GENERATE_PATTERNS) and not has_code:
        return CodeTaskType.GENERATE

    if contains_pattern(normalized_message, EXPLAIN_PATTERNS) and not is_modification_request(normalized_message):
        return CodeTaskType.EXPLAIN

    if has_code and not has_error:
        if is_modification_request(normalized_message):
            return CodeTaskType.REFACTOR
        return CodeTaskType.EXPLAIN

    if is_modification_request(normalized_message) and not has_code:
        return CodeTaskType.GENERATE

    return CodeTaskType.UNKNOWN