#!/usr/bin/env python3
"""
Semantic Correction Classifier

Uses LLM-based classification to detect corrections with high accuracy.
Distinguishes between:
- CORRECTION: User indicating something was WRONG
- PREFERENCE: User wants something different (not wrong)
- ITERATION: Exploring options
- NEW_INFO: Information you couldn't have known
- ADDITION: Extending, not fixing
- CLARIFICATION: Resolving ambiguity
"""

import json
import re
import subprocess
import sys
from dataclasses import dataclass
from typing import Optional

# Classification prompt for the LLM
CLASSIFICATION_PROMPT = """# Correction Detection Classifier

You are analyzing a conversation between a user and an AI assistant (Aria).
Determine if the user's message contains a CORRECTION.

## Definitions

**CORRECTION** = User indicating the assistant was WRONG about something
- Violated a standard, rule, or expectation
- Made a factual error
- Did something incorrectly
- Misunderstood instructions
- Applied a pattern inappropriately
- Forgot something previously discussed

**NOT A CORRECTION:**
- **Preference**: "I'd prefer blue" (not wrong, just different taste)
- **Iteration**: "Let's try a different approach" (exploring options)
- **New info**: "Oh btw, the deadline moved" (info you couldn't have known)
- **Addition**: "Also add X" (extending, not fixing)
- **Clarification**: "I meant the other file" (ambiguity, not error)

## Key Question

> "Did Aria violate something she SHOULD have known or done correctly?"

If yes → CORRECTION
If no → NOT CORRECTION

## Examples

### Example 1: CORRECTION (95%)
**Context**: Aria generated a dark mode landing page
**Message**: "Not all products need dark mode. It should match the brand personality."
**Why**: Aria defaulted to a pattern she shouldn't have. Violated design principle.

### Example 2: CORRECTION (85%)
**Context**: Aria used npm for a new project
**Message**: "We use pnpm for new projects now"
**Why**: Aria forgot an established standard.

### Example 3: NOT CORRECTION (15%)
**Context**: Aria proposed a marketing strategy
**Message**: "Actually the PH market works differently - trust matters more than features"
**Why**: New information Aria couldn't have known. Teaching, not correcting.

### Example 4: NOT CORRECTION (20%)
**Context**: Aria made the header blue
**Message**: "Can we try it in green instead?"
**Why**: Preference/iteration, not an error.

### Example 5: CORRECTION (80%)
**Context**: Aria said the meeting is at 3pm
**Message**: "That's wrong, it's at 2pm"
**Why**: Factual error.

### Example 6: NOT CORRECTION (25%)
**Context**: Aria built feature A
**Message**: "Let's also add feature B"
**Why**: Addition, not fixing something wrong.

### Example 7: CORRECTION (88%)
**Context**: Aria used counting animations on multiple landing pages
**Message**: "The counting up numbers look very repetitive across all the websites"
**Why**: Aria reused patterns inappropriately - violated uniqueness principle.

### Example 8: CORRECTION (85%)
**Context**: Aria delivered a landing page design
**Message**: "Bland, no drama, nothing captivating, feels like stock photo"
**Why**: Output quality fell below acceptable standard.

### Example 9: NOT CORRECTION (30%)
**Context**: Discussing project timeline
**Message**: "Actually, let me think about this more"
**Why**: User pausing to reconsider, not correcting an error.

### Example 10: CORRECTION (90%)
**Context**: Aria shared credentials in chat
**Message**: "Never print API keys in chat messages"
**Why**: Violated security policy.

## Output Format

Return ONLY valid JSON (no markdown, no explanation outside JSON):
{
  "is_correction": boolean,
  "confidence": 0-100,
  "category": "correction" | "preference" | "iteration" | "new_info" | "addition" | "clarification",
  "reasoning": "One sentence explanation",
  "what_was_wrong": "If correction: what specifically was wrong (null if not correction)",
  "scope": "specific" | "general"
}

"scope" meanings:
- "specific": Only applies to this instance
- "general": Applies broadly (pattern, rule, standard)"""


@dataclass
class ClassificationResult:
    is_correction: bool
    confidence: int
    category: str
    reasoning: str
    what_was_wrong: Optional[str]
    scope: Optional[str]
    raw_response: str


def classify_message(
    message: str,
    context: str = "",
    use_llm: bool = False,
    model: str = "claude-haiku-4-5"
) -> ClassificationResult:
    """
    Classify a message using weighted signal analysis.
    
    Args:
        message: The user's message to classify
        context: What Aria said/did before (optional but recommended)
        use_llm: If True, try LLM classification first (slower but more accurate)
        model: Model to use if use_llm=True
    
    Returns:
        ClassificationResult with detection details
    """
    
    user_prompt = f"""Analyze this message:

**Context (what Aria did/said before):**
{context if context else "(No context provided)"}

**User's message:**
{message}

Classify whether this is a CORRECTION."""

    # Default: Use fast weighted signal analysis
    if not use_llm:
        result = fallback_classify(user_prompt)
        try:
            parsed = json.loads(result)
            return ClassificationResult(
                is_correction=parsed.get("is_correction", False),
                confidence=parsed.get("confidence", 0),
                category=parsed.get("category", "unknown"),
                reasoning=parsed.get("reasoning", ""),
                what_was_wrong=parsed.get("what_was_wrong"),
                scope=parsed.get("scope"),
                raw_response=result
            )
        except json.JSONDecodeError:
            pass

    # Optional: Call LLM via CLI for higher accuracy
    try:
        result = call_llm(CLASSIFICATION_PROMPT, user_prompt, model)
        parsed = json.loads(result)
        
        return ClassificationResult(
            is_correction=parsed.get("is_correction", False),
            confidence=parsed.get("confidence", 0),
            category=parsed.get("category", "unknown"),
            reasoning=parsed.get("reasoning", ""),
            what_was_wrong=parsed.get("what_was_wrong"),
            scope=parsed.get("scope"),
            raw_response=result
        )
        
    except json.JSONDecodeError as e:
        return ClassificationResult(
            is_correction=False,
            confidence=0,
            category="error",
            reasoning=f"Failed to parse LLM response: {e}",
            what_was_wrong=None,
            scope=None,
            raw_response=result if 'result' in dir() else ""
        )
    except Exception as e:
        # Final fallback
        result = fallback_classify(user_prompt)
        parsed = json.loads(result)
        return ClassificationResult(
            is_correction=parsed.get("is_correction", False),
            confidence=parsed.get("confidence", 0),
            category=parsed.get("category", "unknown"),
            reasoning=parsed.get("reasoning", ""),
            what_was_wrong=parsed.get("what_was_wrong"),
            scope=parsed.get("scope"),
            raw_response=result
        )


def call_llm(system_prompt: str, user_prompt: str, model: str) -> str:
    """
    Call an LLM with the given prompts.
    Uses the Gemini CLI as it's lightweight and available.
    Falls back to a simple heuristic if no LLM available.
    """
    
    # Try Gemini CLI first (fast, cheap)
    try:
        combined_prompt = f"{system_prompt}\n\n---\n\n{user_prompt}"
        result = subprocess.run(
            ["gemini", "-p", combined_prompt],
            capture_output=True,
            text=True,
            timeout=30
        )
        if result.returncode == 0 and result.stdout.strip():
            # Extract JSON from response (might have markdown wrapping)
            output = result.stdout.strip()
            # Try to find JSON in the output
            if "{" in output:
                start = output.index("{")
                end = output.rindex("}") + 1
                return output[start:end]
            return output
    except (subprocess.TimeoutExpired, FileNotFoundError):
        pass
    
    # Fallback: use regex-based heuristic (less accurate but always works)
    return fallback_classify(user_prompt)


def fallback_classify(user_prompt: str) -> str:
    """
    Weighted signal classification when LLM unavailable.
    Uses the scoring system designed with loewe for accurate detection.
    """
    # Extract message from prompt format
    if "User's message:" in user_prompt:
        message = user_prompt.split("User's message:")[-1].strip()
        # Clean up markdown formatting
        message = message.strip("*").strip()
        # Remove instruction line if present
        if "\n\nClassify" in message:
            message = message.split("\n\nClassify")[0].strip()
        elif "Classify whether" in message:
            message = message.split("Classify whether")[0].strip()
    else:
        message = user_prompt
    
    message_lower = message.lower()
    
    # Base confidence
    confidence = 15
    category = "unknown"
    reasoning_parts = []
    
    # ═══════════════════════════════════════════
    # SUPER-STRONG SIGNALS (+45% each) - These alone indicate correction
    # ═══════════════════════════════════════════
    
    super_strong_patterns = [
        r"that'?s?\s+wrong",  # "that's wrong"
        r"^wrong\b",  # starts with "wrong"
        r"\bincorrect\b",  # "incorrect"
        r"not\s+(?:right|correct)",  # "not right/correct"
        r"you\s+should(?:n'?t|\s+have)",  # "you shouldn't" / "you should have"
    ]
    for pattern in super_strong_patterns:
        if re.search(pattern, message_lower):
            confidence += 55  # Strong enough alone to trigger logging
            reasoning_parts.append("super-strong negation")
            break
    
    # ═══════════════════════════════════════════
    # STRONG SIGNALS (+25% each)
    # ═══════════════════════════════════════════
    
    # Explicit negation (softer forms)
    if "super-strong negation" not in reasoning_parts:
        negation_patterns = [
            r"no[,.]?\s+you", r"that'?s?\s+not",
            r"^no[,.\s]", r"^not\s+all",
            r"very\s+repetitive",  # Strong criticism
        ]
        for pattern in negation_patterns:
            if re.search(pattern, message_lower):
                confidence += 30
                reasoning_parts.append("explicit negation")
                break
    
    # Quality criticism (strong signal - implies work was below standard)
    # Multiple criticisms = stronger signal, so we count them
    quality_criticism = [
        r"\bbland\b", r"no\s+drama", r"nothing\s+captivating",
        r"feels?\s+like\s+(?:stock|template|ai|generic)",
        r"too\s+(?:generic|basic|simple|boring)",
        r"looks?\s+(?:generic|ai|template)",
    ]
    quality_hits = 0
    for pattern in quality_criticism:
        if re.search(pattern, message_lower):
            quality_hits += 1
    if quality_hits > 0:
        # First hit +30, additional hits +15 each
        confidence += 30 + (quality_hits - 1) * 15
        reasoning_parts.append(f"quality criticism (x{quality_hits})")
    
    # Counterfactual ("should have", "shouldn't", "should match", "it should")
    if re.search(r"should(?:n'?t|\s+have|\s+be(?:en)?|\s+not|\s+match|\s+fit|\s+use)", message_lower):
        confidence += 25
        reasoning_parts.append("counterfactual")
    elif re.search(r"it\s+should|you\s+should", message_lower):
        confidence += 25
        reasoning_parts.append("counterfactual")
    
    # Pattern call-out ("always", "every time", "repetitive", "all X")
    if re.search(r"(?:always|every\s+time|repetitive|across\s+all|very\s+repetitive|all\s+(?:the\s+)?(?:products?|sites?|pages?|websites?))", message_lower):
        confidence += 30
        reasoning_parts.append("pattern call-out")
    
    # Memory request ("remember", "don't forget")
    if re.search(r"(?:remember|don'?t\s+forget|keep\s+in\s+mind)", message_lower):
        confidence += 25
        reasoning_parts.append("memory request")
    
    # Instructional correction ("Always use X" + explanation = implicit you weren't doing it)
    # Pattern: "Always [verb]" at start of message + technical explanation
    if re.search(r"^always\s+(?:use|do|run|call|set|specify|include|add)", message_lower):
        # Check if there's an explanation (why the default/current way is wrong)
        if re.search(r"(?:the\s+default|doesn'?t\s+work|needs?\s+to|won'?t\s+work|otherwise|automatically)", message_lower):
            confidence += 35
            reasoning_parts.append("instructional correction")
    
    # ═══════════════════════════════════════════
    # MEDIUM SIGNALS (+15% each)
    # ═══════════════════════════════════════════
    
    # Future imperative ("next time", "from now on")
    if re.search(r"(?:next\s+time|from\s+now|going\s+forward|in\s+the\s+future)", message_lower):
        confidence += 15
        reasoning_parts.append("future imperative")
    
    # Generalization ("all X should", "never do X")
    if re.search(r"(?:all\s+\w+\s+should|never\s+\w+|always\s+\w+)", message_lower):
        confidence += 15
        reasoning_parts.append("generalization")
    
    # External standard reference (strong: implies you violated it)
    if re.search(r"(?:the\s+(?:rule|standard|policy)|we\s+(?:always|never)|supposed\s+to)", message_lower):
        confidence += 25
        reasoning_parts.append("external standard")
    # "We use X for new projects" - implies you should have used it
    elif re.search(r"we\s+use\s+\w+\s+(?:for|now|here)", message_lower):
        confidence += 55  # Very strong - directly states what should be used (implies you didn't)
        reasoning_parts.append("standard reminder")
    
    # Frustration markers
    if re.search(r"(?:why\s+did\s+you|i\s+(?:already|just)\s+(?:told|said)|again\?)", message_lower):
        confidence += 15
        reasoning_parts.append("frustration markers")
    
    # ═══════════════════════════════════════════
    # WEAK SIGNALS (+10% each)
    # ═══════════════════════════════════════════
    
    # Contrast + judgment
    if re.search(r"(?:but\s+(?:it|that|this)\s+(?:is|was|looks)|too\s+\w+|not\s+enough)", message_lower):
        confidence += 10
        reasoning_parts.append("contrast judgment")
    
    # Teaching frame
    if re.search(r"(?:here'?s?\s+how|the\s+(?:correct|right)\s+way|you\s+need\s+to)", message_lower):
        confidence += 10
        reasoning_parts.append("teaching frame")
    
    # Emphasis (ALLCAPS words or multiple exclamations - not just any caps)
    if re.search(r'\b[A-Z]{4,}\b|!{2,}', message):
        confidence += 10
        reasoning_parts.append("emphasis")
    
    # ═══════════════════════════════════════════
    # NEGATIVE SIGNALS (-20% each)
    # ═══════════════════════════════════════════
    
    # Preference language
    if re.search(r"(?:i'?d?\s+prefer|i\s+like|can\s+we\s+try|how\s+about|what\s+if)", message_lower):
        confidence -= 20
        reasoning_parts.append("preference language (negative)")
    
    # Exploratory tone
    if re.search(r"(?:let'?s?\s+try|maybe|perhaps|could\s+we|experiment)", message_lower):
        confidence -= 20
        reasoning_parts.append("exploratory tone (negative)")
    
    # New info framing
    if re.search(r"(?:btw|by\s+the\s+way|fyi|just\s+so\s+you\s+know|oh\s+also)", message_lower):
        confidence -= 20
        reasoning_parts.append("new info framing (negative)")
    
    # Addition framing
    if re.search(r"(?:also\s+add|additionally|and\s+also|plus)", message_lower):
        confidence -= 20
        reasoning_parts.append("addition framing (negative)")
    
    # ═══════════════════════════════════════════
    # DETERMINE CATEGORY
    # ═══════════════════════════════════════════
    
    confidence = max(0, min(100, confidence))
    
    if confidence >= 70:
        category = "correction"
        is_correction = True
    elif confidence >= 40:
        category = "uncertain"
        is_correction = False  # Don't auto-log uncertain
    else:
        # Determine non-correction category
        if "preference language" in str(reasoning_parts):
            category = "preference"
        elif "exploratory tone" in str(reasoning_parts):
            category = "iteration"
        elif "new info framing" in str(reasoning_parts):
            category = "new_info"
        elif "addition framing" in str(reasoning_parts):
            category = "addition"
        else:
            category = "not_correction"
        is_correction = False
    
    reasoning = f"Weighted signals: {', '.join(reasoning_parts)}" if reasoning_parts else "No strong signals detected"
    
    return json.dumps({
        "is_correction": is_correction,
        "confidence": confidence,
        "category": category,
        "reasoning": reasoning,
        "what_was_wrong": None if not is_correction else "See original message",
        "scope": "general" if "pattern call-out" in str(reasoning_parts) or "generalization" in str(reasoning_parts) else "specific"
    })


def classify_thread(
    messages: list[dict],
    aria_responses: list[dict] = None,
    recency_decay: bool = True
) -> list[ClassificationResult]:
    """
    Classify all messages in a thread efficiently.
    
    Args:
        messages: List of {"text": str, "author": str, "timestamp": int}
        aria_responses: Aria's responses for context matching
        recency_decay: Apply decay to older messages
    
    Returns:
        List of ClassificationResult for detected corrections
    """
    results = []
    total_messages = len(messages)
    
    for i, msg in enumerate(messages):
        # Skip Aria's own messages
        if msg.get("author", "").lower() in ["aria", "aria#5762"]:
            continue
        
        # Find context (what Aria said before this message)
        context = ""
        if aria_responses:
            # Find most recent Aria response before this message
            for resp in reversed(aria_responses):
                if resp.get("timestamp", 0) < msg.get("timestamp", 0):
                    context = resp.get("text", "")
                    break
        
        # Classify
        result = classify_message(msg["text"], context)
        
        # Apply recency decay
        if recency_decay and result.is_correction:
            position_from_end = total_messages - i
            decay = get_decay_factor(position_from_end)
            result.confidence = int(result.confidence * decay)
        
        # Only keep if confidence meets threshold
        if result.confidence >= 40:  # Include flagged ones too
            results.append(result)
    
    return results


def get_decay_factor(position_from_end: int) -> float:
    """
    Apply recency decay - older messages less likely to be actionable corrections.
    """
    if position_from_end <= 5:
        return 1.0
    elif position_from_end <= 15:
        return 0.7
    elif position_from_end <= 30:
        return 0.4
    else:
        return 0.2


def main():
    """CLI interface for testing."""
    if len(sys.argv) < 2:
        print("Usage: semantic_classifier.py <message> [context]")
        print("       semantic_classifier.py --test")
        sys.exit(1)
    
    if sys.argv[1] == "--test":
        # Run test cases
        test_cases = [
            ("Not all products need dark mode. It should match the brand personality.", 
             "Aria generated dark mode landing pages for all products"),
            ("Can we try it in green instead?", 
             "Aria made the header blue"),
            ("We use pnpm for new projects now", 
             "Aria used npm to install packages"),
            ("Actually the PH market works differently - trust matters more",
             "Aria proposed a Western marketing strategy"),
            ("Let's also add a contact form",
             "Aria built the landing page"),
        ]
        
        print("Running test cases...\n")
        for msg, ctx in test_cases:
            result = classify_message(msg, ctx)
            print(f"Message: {msg[:50]}...")
            print(f"  → {result.category} ({result.confidence}%)")
            print(f"  → Correction: {result.is_correction}")
            print(f"  → Reasoning: {result.reasoning}")
            print()
        return
    
    message = sys.argv[1]
    context = sys.argv[2] if len(sys.argv) > 2 else ""
    
    result = classify_message(message, context)
    print(json.dumps({
        "is_correction": result.is_correction,
        "confidence": result.confidence,
        "category": result.category,
        "reasoning": result.reasoning,
        "what_was_wrong": result.what_was_wrong,
        "scope": result.scope,
    }, indent=2))


if __name__ == "__main__":
    main()
