"""
BrowseComp Answer Decoding Pipeline

This module handles encoded answers found in BrowseComp datasets.
Some BrowseComp answers appear to be encoded (e.g., "Y00Qh+ep") and need
decoding to extract the actual answer.

Based on BROWSECOMP_IMPROVEMENT_STRATEGY.md recommendations.
"""

import base64
from loguru import logger
import re
import urllib.parse
from typing import Optional, Tuple


class BrowseCompAnswerDecoder:
    """
    Handle encoded BrowseComp answers with multiple decoding schemes.

    Features:
    1. Automatic encoding detection
    2. Multiple decoding scheme support
    3. Answer validation
    4. Fallback to original if decoding fails
    """

    def __init__(self):
        self.encoding_schemes = [
            "base64",
            "hex",
            "url_encoding",
            "rot13",
            "caesar_cipher",
        ]

        # Patterns that suggest encoded content
        self.encoded_patterns = [
            r"^[A-Za-z0-9+/]+=*$",  # Base64 pattern
            r"^[0-9A-Fa-f]+$",  # Hex pattern
            r"%[0-9A-Fa-f]{2}",  # URL encoded
            r"^[A-Za-z0-9]{8,}$",  # Random string pattern
        ]

    def decode_answer(self, raw_answer: str) -> Tuple[str, Optional[str]]:
        """
        Attempt to decode a potentially encoded answer.

        Args:
            raw_answer: The raw answer string that may be encoded

        Returns:
            Tuple of (decoded_answer, encoding_scheme_used)
            If no decoding works, returns (original_answer, None)
        """
        if not raw_answer or len(raw_answer.strip()) == 0:
            return raw_answer, None

        # Clean the input
        clean_answer = raw_answer.strip()

        # Check if answer looks like plaintext first
        if self.is_likely_direct_answer(clean_answer):
            logger.debug(f"Answer appears to be plaintext: {clean_answer}")
            return clean_answer, None

        logger.info(
            f"Attempting to decode potentially encoded answer: {clean_answer}"
        )

        # Try each encoding scheme
        for scheme in self.encoding_schemes:
            try:
                decoded = self.apply_decoding_scheme(clean_answer, scheme)
                if decoded and self.validate_decoded_answer(decoded):
                    logger.info(
                        f"Successfully decoded using {scheme}: {clean_answer} -> {decoded}"
                    )
                    return decoded, scheme

            except Exception as e:
                logger.debug(f"Failed to decode with {scheme}: {e}")
                continue

        # No successful decoding
        logger.warning(
            f"Could not decode answer, returning original: {clean_answer}"
        )
        return clean_answer, None

    def is_likely_direct_answer(self, answer: str) -> bool:
        """
        Check if answer looks like plaintext rather than encoded.

        Args:
            answer: The answer string to check

        Returns:
            True if answer appears to be plaintext
        """
        # Very short answers are likely plaintext
        if len(answer) < 4:
            return True

        # Check for common English words
        english_indicators = [
            "the",
            "and",
            "or",
            "of",
            "in",
            "to",
            "a",
            "an",
            "company",
            "group",
            "inc",
            "ltd",
            "corp",
            "corporation",
            "person",
            "people",
            "event",
            "year",
            "years",
            "million",
            "billion",
            "thousand",
        ]

        answer_lower = answer.lower()
        if any(word in answer_lower for word in english_indicators):
            return True

        # Check for sentence-like structure
        if " " in answer and len(answer.split()) > 1:
            # Has spaces and multiple words - likely plaintext
            return True

        # Check if it matches common answer patterns
        common_patterns = [
            r"^\d{4}$",  # Year
            r"^\$?\d+\.?\d*[KMB]?$",  # Number/money
            r"^[A-Z][a-z]+ [A-Z][a-z]+$",  # Name format
            r"^\d+%$",  # Percentage
        ]

        for pattern in common_patterns:
            if re.match(pattern, answer):
                return True

        # Check character distribution - encoded text often has unusual distribution
        char_diversity = (
            len(set(answer)) / len(answer) if len(answer) > 0 else 0
        )
        if char_diversity < 0.3:  # Low diversity suggests repetitive/encoded
            return False

        # If none of the encoded patterns match, probably plaintext
        is_encoded = any(
            re.search(pattern, answer) for pattern in self.encoded_patterns
        )
        return not is_encoded

    def apply_decoding_scheme(self, text: str, scheme: str) -> Optional[str]:
        """
        Apply a specific decoding scheme to text.

        Args:
            text: Text to decode
            scheme: Decoding scheme to use

        Returns:
            Decoded text or None if decoding fails
        """
        try:
            if scheme == "base64":
                return self._decode_base64(text)
            elif scheme == "hex":
                return self._decode_hex(text)
            elif scheme == "url_encoding":
                return self._decode_url(text)
            elif scheme == "rot13":
                return self._decode_rot13(text)
            elif scheme == "caesar_cipher":
                return self._decode_caesar(text)
            else:
                logger.warning(f"Unknown decoding scheme: {scheme}")
                return None

        except Exception as e:
            logger.debug(f"Failed to apply {scheme} decoding: {e}")
            return None

    def _decode_base64(self, text: str) -> Optional[str]:
        """Decode base64 encoded text."""
        try:
            # Add padding if needed
            missing_padding = len(text) % 4
            if missing_padding:
                text += "=" * (4 - missing_padding)

            decoded_bytes = base64.b64decode(text)
            return decoded_bytes.decode("utf-8")

        except Exception:
            return None

    def _decode_hex(self, text: str) -> Optional[str]:
        """Decode hexadecimal encoded text."""
        try:
            # Remove any whitespace or non-hex characters
            clean_hex = re.sub(r"[^0-9A-Fa-f]", "", text)

            # Must have even length
            if len(clean_hex) % 2 != 0:
                return None

            decoded_bytes = bytes.fromhex(clean_hex)
            return decoded_bytes.decode("utf-8")

        except Exception:
            return None

    def _decode_url(self, text: str) -> Optional[str]:
        """Decode URL encoded text."""
        try:
            return urllib.parse.unquote(text)
        except Exception:
            return None

    def _decode_rot13(self, text: str) -> Optional[str]:
        """Decode ROT13 encoded text."""
        try:
            import codecs

            return codecs.decode(text, "rot13")
        except Exception:
            return None

    def _decode_caesar(self, text: str) -> Optional[str]:
        """
        Try different Caesar cipher shifts.
        Returns the most English-like result.
        """
        best_result = None
        best_score = 0

        # Try shifts 1-25
        for shift in range(1, 26):
            try:
                decoded = self._caesar_shift(text, shift)
                score = self._english_score(decoded)

                if score > best_score:
                    best_score = score
                    best_result = decoded

            except Exception:
                continue

        # Only return if it looks reasonably English-like
        return best_result if best_score > 0.3 else None

    def _caesar_shift(self, text: str, shift: int) -> str:
        """Apply Caesar cipher shift."""
        result = []

        for char in text:
            if char.isalpha():
                # Determine if uppercase or lowercase
                start = ord("A") if char.isupper() else ord("a")
                # Apply shift with wraparound
                shifted = (ord(char) - start + shift) % 26 + start
                result.append(chr(shifted))
            else:
                result.append(char)

        return "".join(result)

    def _english_score(self, text: str) -> float:
        """
        Score how English-like a text appears.
        Simple heuristic based on common letters and words.
        """
        if not text:
            return 0.0

        text_lower = text.lower()

        # Common English letter frequencies (approximate)
        common_letters = "etaoinshrdlcumwfgypbvkjxqz"
        letter_score = 0
        letter_count = 0

        for char in text_lower:
            if char.isalpha():
                letter_count += 1
                # More common letters get higher scores
                if char in common_letters[:10]:  # Top 10 most common
                    letter_score += 2
                elif char in common_letters[:20]:  # Top 20
                    letter_score += 1

        if letter_count == 0:
            return 0.0

        base_score = letter_score / letter_count

        # Bonus for common English words
        common_words = [
            "the",
            "and",
            "of",
            "to",
            "a",
            "in",
            "is",
            "it",
            "you",
            "that",
        ]
        word_bonus = sum(1 for word in common_words if word in text_lower)

        return min(1.0, base_score + word_bonus * 0.1)

    def validate_decoded_answer(self, decoded: str) -> bool:
        """
        Validate that decoded text looks like a reasonable answer.

        Args:
            decoded: The decoded text to validate

        Returns:
            True if decoded text appears valid
        """
        if not decoded or len(decoded.strip()) == 0:
            return False

        # Remove leading/trailing whitespace
        decoded = decoded.strip()

        # Check length - should be reasonable
        if len(decoded) < 1 or len(decoded) > 1000:
            return False

        # Check for readable characters
        printable_count = sum(1 for c in decoded if c.isprintable())
        if printable_count / len(decoded) < 0.8:  # At least 80% printable
            return False

        # Check for control characters (bad sign)
        if any(ord(c) < 32 and c not in "\t\n\r" for c in decoded):
            return False

        # Check character distribution
        char_types = {
            "alpha": sum(1 for c in decoded if c.isalpha()),
            "digit": sum(1 for c in decoded if c.isdigit()),
            "space": sum(1 for c in decoded if c.isspace()),
            "punct": sum(
                1 for c in decoded if not c.isalnum() and not c.isspace()
            ),
        }

        total_chars = len(decoded)

        # Should have some letters
        if char_types["alpha"] / total_chars < 0.3:
            return False

        # Shouldn't be mostly punctuation
        if char_types["punct"] / total_chars > 0.5:
            return False

        return True

    def analyze_answer_encoding(self, answer: str) -> dict:
        """
        Analyze an answer to determine likely encoding type.

        Returns analysis results for debugging/logging.
        """
        analysis = {
            "original": answer,
            "length": len(answer),
            "likely_plaintext": self.is_likely_direct_answer(answer),
            "pattern_matches": [],
            "attempted_decodings": {},
        }

        # Check which patterns match
        for i, pattern in enumerate(self.encoded_patterns):
            if re.search(pattern, answer):
                analysis["pattern_matches"].append(
                    {
                        "pattern": pattern,
                        "type": ["base64", "hex", "url", "random"][i],
                    }
                )

        # Try each decoding scheme
        for scheme in self.encoding_schemes:
            try:
                decoded = self.apply_decoding_scheme(answer, scheme)
                is_valid = (
                    self.validate_decoded_answer(decoded) if decoded else False
                )

                analysis["attempted_decodings"][scheme] = {
                    "decoded": decoded,
                    "valid": is_valid,
                    "length": len(decoded) if decoded else 0,
                }
            except Exception:
                logger.exception(f"Error decoding with scheme {scheme}")
                analysis["attempted_decodings"][scheme] = {
                    "error": "Decoding failed"
                }

        return analysis
