"""Few-shot examples for prompt engineering.

Curated examples for different error types to improve
LLM fix generation quality.
"""

from typing import Dict, List

from autodebug.prompts.examples.assertion_error import (
    ASSERTION_ERROR_EXAMPLES,
    ASSERTION_ERROR_NEGATIVE_EXAMPLES,
)
from autodebug.prompts.examples.attribute_error import (
    ATTRIBUTE_ERROR_EXAMPLES,
    ATTRIBUTE_ERROR_NEGATIVE_EXAMPLES,
)
from autodebug.prompts.examples.generic import (
    GENERIC_EXAMPLES,
    GENERIC_NEGATIVE_EXAMPLES,
)
from autodebug.prompts.examples.key_error import (
    KEY_ERROR_EXAMPLES,
    KEY_ERROR_NEGATIVE_EXAMPLES,
)
from autodebug.prompts.examples.timeout_error import (
    TIMEOUT_ERROR_EXAMPLES,
    TIMEOUT_ERROR_NEGATIVE_EXAMPLES,
)
from autodebug.prompts.examples.type_error import (
    TYPE_ERROR_EXAMPLES,
    TYPE_ERROR_NEGATIVE_EXAMPLES,
)
from autodebug.prompts.examples.value_error import (
    VALUE_ERROR_EXAMPLES,
    VALUE_ERROR_NEGATIVE_EXAMPLES,
)

# Map error types to their examples
EXAMPLES_BY_TYPE: Dict[str, List[Dict]] = {
    "KeyError": KEY_ERROR_EXAMPLES,
    "TypeError": TYPE_ERROR_EXAMPLES,
    "AttributeError": ATTRIBUTE_ERROR_EXAMPLES,
    "ValueError": VALUE_ERROR_EXAMPLES,
    "AssertionError": ASSERTION_ERROR_EXAMPLES,
    "AssertionError": ASSERTION_ERROR_EXAMPLES,
    "TimeoutError": TIMEOUT_ERROR_EXAMPLES,
    "asyncio.TimeoutError": TIMEOUT_ERROR_EXAMPLES,
}

NEGATIVE_EXAMPLES_BY_TYPE: Dict[str, List[Dict]] = {
    "KeyError": KEY_ERROR_NEGATIVE_EXAMPLES,
    "TypeError": TYPE_ERROR_NEGATIVE_EXAMPLES,
    "AttributeError": ATTRIBUTE_ERROR_NEGATIVE_EXAMPLES,
    "ValueError": VALUE_ERROR_NEGATIVE_EXAMPLES,
    "AssertionError": ASSERTION_ERROR_NEGATIVE_EXAMPLES,
    "AssertionError": ASSERTION_ERROR_NEGATIVE_EXAMPLES,
    "TimeoutError": TIMEOUT_ERROR_NEGATIVE_EXAMPLES,
    "asyncio.TimeoutError": TIMEOUT_ERROR_NEGATIVE_EXAMPLES,
}


def load_examples_for_type(error_type: str, limit: int = 3) -> List[Dict]:
    """Load few-shot examples for an error type.

    Args:
        error_type: The exception type (e.g., "KeyError")
        limit: Maximum number of examples to return

    Returns:
        List of example dictionaries
    """
    # Try exact match first
    if error_type in EXAMPLES_BY_TYPE:
        return EXAMPLES_BY_TYPE[error_type][:limit]

    # Try suffix match (e.g., "requests.exceptions.ConnectionError" -> "ConnectionError")
    base_type = error_type.split(".")[-1]
    if base_type in EXAMPLES_BY_TYPE:
        return EXAMPLES_BY_TYPE[base_type][:limit]

    # Fall back to generic examples
    return GENERIC_EXAMPLES[:limit]


def load_negative_examples_for_type(error_type: str, limit: int = 2) -> List[Dict]:
    """Load negative examples (what NOT to do) for an error type.

    Args:
        error_type: The exception type
        limit: Maximum number of examples to return

    Returns:
        List of negative example dictionaries
    """
    # Try exact match first
    if error_type in NEGATIVE_EXAMPLES_BY_TYPE:
        return NEGATIVE_EXAMPLES_BY_TYPE[error_type][:limit]

    # Try suffix match
    base_type = error_type.split(".")[-1]
    if base_type in NEGATIVE_EXAMPLES_BY_TYPE:
        return NEGATIVE_EXAMPLES_BY_TYPE[base_type][:limit]

    # Fall back to generic negative examples
    return GENERIC_NEGATIVE_EXAMPLES[:limit]


__all__ = [
    "load_examples_for_type",
    "load_negative_examples_for_type",
    "EXAMPLES_BY_TYPE",
    "NEGATIVE_EXAMPLES_BY_TYPE",
]
