from __future__ import annotations

import re
from enum import Enum

from epochler.config import Settings


class ErrorCategory(str, Enum):
    SYNTAX = "syntax"
    API_MISUSE = "api_misuse"
    IMPORT = "import_error"
    DATA = "data"
    RUNTIME = "runtime"
    LOGIC = "logic"
    UNKNOWN = "unknown"


TIER1_CATEGORIES = {
    ErrorCategory.SYNTAX,
    ErrorCategory.API_MISUSE,
    ErrorCategory.IMPORT,
    ErrorCategory.DATA,
    ErrorCategory.UNKNOWN,
}

TIER2_CATEGORIES = {
    ErrorCategory.RUNTIME,
    ErrorCategory.LOGIC,
}

_PATTERNS: list[tuple[re.Pattern, ErrorCategory]] = [
    (re.compile(r"SyntaxError", re.IGNORECASE), ErrorCategory.SYNTAX),
    (re.compile(r"IndentationError", re.IGNORECASE), ErrorCategory.SYNTAX),
    (re.compile(r"TabError", re.IGNORECASE), ErrorCategory.SYNTAX),

    (re.compile(r"ModuleNotFoundError", re.IGNORECASE), ErrorCategory.IMPORT),
    (re.compile(r"ImportError", re.IGNORECASE), ErrorCategory.IMPORT),
    (re.compile(r"No module named", re.IGNORECASE), ErrorCategory.IMPORT),

    (re.compile(r"TypeError.*got an unexpected keyword argument"), ErrorCategory.API_MISUSE),
    (re.compile(r"TypeError.*takes \d+ positional argument"), ErrorCategory.API_MISUSE),
    (re.compile(r"TypeError.*missing \d+ required"), ErrorCategory.API_MISUSE),
    (re.compile(r"TypeError.*argument"), ErrorCategory.API_MISUSE),
    (re.compile(r"AttributeError"), ErrorCategory.API_MISUSE),
    (re.compile(r"DeprecationWarning"), ErrorCategory.API_MISUSE),

    (re.compile(r"FileNotFoundError"), ErrorCategory.DATA),
    (re.compile(r"KeyError"), ErrorCategory.DATA),
    (re.compile(r"ValueError.*shape"), ErrorCategory.DATA),
    (re.compile(r"ParserError"), ErrorCategory.DATA),
    (re.compile(r"EmptyDataError"), ErrorCategory.DATA),
    (re.compile(r"No such file or directory"), ErrorCategory.DATA),
    (re.compile(r"not found.*column", re.IGNORECASE), ErrorCategory.DATA),
    (re.compile(r"column.*not found", re.IGNORECASE), ErrorCategory.DATA),

    (re.compile(r"MemoryError"), ErrorCategory.RUNTIME),
    (re.compile(r"CUDA out of memory", re.IGNORECASE), ErrorCategory.RUNTIME),
    (re.compile(r"OutOfMemoryError", re.IGNORECASE), ErrorCategory.RUNTIME),
    (re.compile(r"RuntimeError.*out of memory", re.IGNORECASE), ErrorCategory.RUNTIME),
    (re.compile(r"Killed", re.IGNORECASE), ErrorCategory.RUNTIME),

    (re.compile(r"\bnan\b", re.IGNORECASE), ErrorCategory.LOGIC),
    (re.compile(r"\binf\b"), ErrorCategory.LOGIC),
    (re.compile(r"convergence", re.IGNORECASE), ErrorCategory.LOGIC),
    (re.compile(r"diverge", re.IGNORECASE), ErrorCategory.LOGIC),
    (re.compile(r"loss.*nan", re.IGNORECASE), ErrorCategory.LOGIC),
]


def classify_error(stderr: str, exit_code: int = 1, *, timed_out: bool = False) -> ErrorCategory:
    """Classify a script error based on stderr content."""
    if timed_out:
        return ErrorCategory.RUNTIME

    for pattern, category in _PATTERNS:
        if pattern.search(stderr):
            return category

    return ErrorCategory.UNKNOWN


def select_model(
    category: ErrorCategory,
    fix_attempts: int,
    settings: Settings,
) -> str:
    """Pick the right LLM model based on error category and retry count.

    Tier 1 errors (syntax, API misuse, imports, data, unknown) start with
    the cheaper/faster planner model (Sonnet) and escalate to the full model
    (Opus) after 2 failed fix attempts.

    Tier 2 errors (runtime, logic) go straight to the full model.
    """
    if category in TIER2_CATEGORIES:
        return settings.anthropic_model
    if fix_attempts >= 2:
        return settings.anthropic_model
    return settings.anthropic_planner_model
