import re
from typing import Any, Dict, List

def classify_message(message: str) -> str:
    """Classify message type using LangChain approach."""
    message_lower = message.lower().strip()
    
    # Programming-related keywords
    code_keywords = [
        "write code", "python", "function", "class", "algorithm", 
        "programming", "code", "script", "loop", "variable", "import",
        "def ", "return", "if ", "for ", "while ", "try:", "except:",
        "list", "dict", "array", "dataframe", "numpy", "pandas"
    ]
    
    # Conversational keywords
    conversation_keywords = [
        "hi", "hello", "how are you", "what's up", "good morning",
        "good afternoon", "good evening", "bye", "goodbye", "thank you",
        "thanks", "help", "who are you", "what can you do"
    ]
    
    # Check for code-related content
    if any(keyword in message_lower for keyword in code_keywords):
        return "code"
    
    # Check for conversational content
    if any(keyword in message_lower for keyword in conversation_keywords):
        return "conversation"
    
    # Check for question marks or short messages
    if message_lower.endswith("?") or len(message_lower.split()) <= 5:
        return "conversation"
    
    # Default to conversation for ambiguous cases
    return "conversation"


def process_response(raw_response: str, response_type: str) -> Dict[str, Any]:
    """Process and format the model's response."""
    if response_type == "conversation":
        return {"response": raw_response.strip()}
    elif response_type == "code":
        code_match = re.search(r"```python\n(.*?)\n```", raw_response, re.DOTALL)
        if code_match:
            return {"generated_code": code_match.group(1).strip()}
        return {"generated_code": raw_response.strip()}
    elif response_type == "explanation":
        explanation = re.sub(r"```python\n.*?\n```", "", raw_response, flags=re.DOTALL).strip()
        return {"explanation": explanation}
    else:  # "both"
        code = None
        explanation = raw_response
        code_match = re.search(r"```python\n(.*?)\n```", raw_response, re.DOTALL)
        if code_match:
            code = code_match.group(1).strip()
            explanation = re.sub(r"```python\n.*?\n```", "", raw_response, flags=re.DOTALL).strip()
        explanation = re.sub(r"\[EXPLANATION\]\s*", "", explanation, flags=re.IGNORECASE).strip()
        explanation = re.sub(r"\[CODE\]\s*", "", explanation, flags=re.IGNORECASE).strip()
        return {
            "response": raw_response,
            "generated_code": code,
            "explanation": explanation
        }