"""JSON structure handler.

Extracts structural elements from JSON content:
- Keys (navigational - tells LLM what fields exist)
- Brackets and colons (structural syntax)
- Short values like booleans, nulls, small numbers

Values (strings, long numbers, nested content) are marked as compressible.

This enables the LLM to see the full schema while values are compressed.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from enum import Enum
from typing import Any, cast

from headroom.compression.handlers.base import BaseStructureHandler, HandlerResult
from headroom.compression.masks import EntropyScore, StructureMask


class JSONTokenType(Enum):
    """Types of JSON tokens for structure detection."""

    KEY = "key"  # Object key (always structural)
    STRING_VALUE = "string_value"  # String value (compressible)
    NUMBER = "number"  # Numeric value (preserve if short)
    BOOLEAN = "boolean"  # true/false (always structural)
    NULL = "null"  # null (always structural)
    BRACKET = "bracket"  # {, }, [, ] (always structural)
    COLON = "colon"  # : (always structural)
    COMMA = "comma"  # , (always structural)
    WHITESPACE = "whitespace"  # spaces, newlines (compressible)


@dataclass
class JSONToken:
    """A token in JSON content with its type and position."""

    text: str
    token_type: JSONTokenType
    start: int
    end: int

    @property
    def is_structural(self) -> bool:
        """Whether this token should be preserved."""
        return self.token_type in (
            JSONTokenType.KEY,
            JSONTokenType.BOOLEAN,
            JSONTokenType.NULL,
            JSONTokenType.BRACKET,
            JSONTokenType.COLON,
            JSONTokenType.COMMA,
        )


class JSONStructureHandler(BaseStructureHandler):
    """Handler for JSON content.

    Preserves:
    - All keys (navigational - LLM sees what fields exist)
    - Structural syntax ({, }, [, ], :, ,)
    - Booleans and nulls (small, semantically important)
    - High-entropy strings (UUIDs, hashes - identifiers)
    - Short numbers (often IDs or important values)

    Compresses:
    - Long string values (descriptions, content)
    - Whitespace
    - Redundant array elements (after first few)

    Example:
        >>> handler = JSONStructureHandler()
        >>> result = handler.get_mask('{"name": "Alice", "id": "usr_123"}')
        >>> # Keys "name" and "id" preserved, values may be compressed
    """

    def __init__(
        self,
        preserve_short_values: bool = True,
        short_value_threshold: int = 20,
        preserve_high_entropy: bool = True,
        entropy_threshold: float = 0.85,
        max_array_items_full: int = 3,  # Keep first N items fully
        max_number_digits: int = 10,  # Preserve numbers up to N digits
    ):
        """Initialize the JSON handler.

        Args:
            preserve_short_values: Preserve short string values.
            short_value_threshold: Max length for "short" values.
            preserve_high_entropy: Preserve high-entropy strings (UUIDs, etc.).
            entropy_threshold: Entropy threshold for preservation.
            max_array_items_full: Number of array items to keep in full.
            max_number_digits: Max digits for numbers to preserve (often IDs).
        """
        super().__init__(name="json")
        self.preserve_short_values = preserve_short_values
        self.short_value_threshold = short_value_threshold
        self.preserve_high_entropy = preserve_high_entropy
        self.entropy_threshold = entropy_threshold
        self.max_array_items_full = max_array_items_full
        self.max_number_digits = max_number_digits

    def can_handle(self, content: str) -> bool:
        """Check if content is valid JSON."""
        stripped = content.strip()
        if not stripped.startswith(("{", "[")):
            return False
        try:
            json.loads(stripped)
            return True
        except (json.JSONDecodeError, ValueError):
            return False

    def _extract_mask(
        self,
        content: str,
        tokens: list[str],
        **kwargs: Any,
    ) -> HandlerResult:
        """Extract structure mask from JSON content.

        Args:
            content: JSON content.
            tokens: Character-level tokens.
            **kwargs: Additional options.

        Returns:
            HandlerResult with mask marking structural elements.
        """
        # Tokenize JSON to identify structure
        json_tokens = self._tokenize_json(content)

        # Build character-level mask
        mask = [False] * len(content)

        # Track array depth for selective preservation
        array_depth = 0
        array_item_counts: dict[int, int] = {}  # depth -> count

        for token in json_tokens:
            # Track array items
            if token.token_type == JSONTokenType.BRACKET:
                if token.text == "[":
                    array_depth += 1
                    array_item_counts[array_depth] = 0
                elif token.text == "]":
                    if array_depth in array_item_counts:
                        del array_item_counts[array_depth]
                    array_depth = max(0, array_depth - 1)

            # Count array items at commas
            if token.token_type == JSONTokenType.COMMA and array_depth > 0:
                array_item_counts[array_depth] = array_item_counts.get(array_depth, 0) + 1

            # Determine if this token should be preserved
            preserve = self._should_preserve_token(
                token,
                array_depth,
                array_item_counts.get(array_depth, 0),
            )

            # Mark in mask
            if preserve:
                for i in range(token.start, min(token.end, len(mask))):
                    mask[i] = True

        # Convert to character tokens if needed
        char_tokens = list(content) if tokens == list(content) else tokens

        return HandlerResult(
            mask=StructureMask(tokens=char_tokens, mask=mask),
            handler_name=self.name,
            confidence=1.0,
            metadata={
                "token_count": len(json_tokens),
                "key_count": sum(1 for t in json_tokens if t.token_type == JSONTokenType.KEY),
            },
        )

    def _should_preserve_token(
        self,
        token: JSONToken,
        array_depth: int,
        array_item_index: int,
    ) -> bool:
        """Determine if a token should be preserved.

        Args:
            token: The JSON token.
            array_depth: Current array nesting depth.
            array_item_index: Index of current item in array.

        Returns:
            True if token should be preserved.
        """
        # Always preserve structural tokens
        if token.is_structural:
            return True

        # Whitespace is never preserved
        if token.token_type == JSONTokenType.WHITESPACE:
            return False

        # Numbers: preserve short ones (often IDs)
        if token.token_type == JSONTokenType.NUMBER:
            return len(token.text) <= self.max_number_digits

        # String values: selective preservation
        if token.token_type == JSONTokenType.STRING_VALUE:
            # Check if we're past the max array items threshold
            if array_depth > 0 and array_item_index >= self.max_array_items_full:
                # In deep array, be more aggressive
                return False

            # Preserve short values
            if self.preserve_short_values and len(token.text) <= self.short_value_threshold:
                return True

            # Preserve high-entropy values (UUIDs, hashes)
            if self.preserve_high_entropy:
                # Strip quotes for entropy calculation
                value = token.text.strip('"')
                score = EntropyScore.compute(value, self.entropy_threshold)
                if score.should_preserve:
                    return True

            return False

        return False

    def _tokenize_json(self, content: str) -> list[JSONToken]:
        """Tokenize JSON content into typed tokens.

        This is a simple tokenizer that identifies JSON structure.
        It's not a full parser - just enough to identify keys vs values.

        Args:
            content: JSON content.

        Returns:
            List of JSONToken objects.
        """
        tokens: list[JSONToken] = []
        i = 0
        n = len(content)

        # Track if we're expecting a key (after { or ,)
        expect_key = False
        brace_stack: list[str] = []

        while i < n:
            char = content[i]

            # Whitespace
            if char in " \t\n\r":
                start = i
                while i < n and content[i] in " \t\n\r":
                    i += 1
                tokens.append(JSONToken(content[start:i], JSONTokenType.WHITESPACE, start, i))
                continue

            # Brackets
            if char in "{}[]":
                tokens.append(JSONToken(char, JSONTokenType.BRACKET, i, i + 1))
                if char == "{":
                    brace_stack.append("{")
                    expect_key = True
                elif char == "}":
                    if brace_stack and brace_stack[-1] == "{":
                        brace_stack.pop()
                    expect_key = False
                elif char == "[":
                    brace_stack.append("[")
                    expect_key = False
                elif char == "]":
                    if brace_stack and brace_stack[-1] == "[":
                        brace_stack.pop()
                i += 1
                continue

            # Colon
            if char == ":":
                tokens.append(JSONToken(char, JSONTokenType.COLON, i, i + 1))
                expect_key = False
                i += 1
                continue

            # Comma
            if char == ",":
                tokens.append(JSONToken(char, JSONTokenType.COMMA, i, i + 1))
                # After comma in object, expect key
                if brace_stack and brace_stack[-1] == "{":
                    expect_key = True
                i += 1
                continue

            # String (key or value)
            if char == '"':
                start = i
                i += 1
                while i < n and content[i] != '"':
                    if content[i] == "\\":
                        i += 2  # Skip escaped character
                    else:
                        i += 1
                i += 1  # Include closing quote

                text = content[start:i]

                # Determine if this is a key or value
                # Look ahead for colon (skipping whitespace)
                j = i
                while j < n and content[j] in " \t\n\r":
                    j += 1

                is_key = j < n and content[j] == ":" and expect_key

                if is_key:
                    tokens.append(JSONToken(text, JSONTokenType.KEY, start, i))
                    expect_key = False
                else:
                    tokens.append(JSONToken(text, JSONTokenType.STRING_VALUE, start, i))

                continue

            # Number
            if char in "-0123456789":
                start = i
                # Match JSON number pattern
                if char == "-":
                    i += 1
                while i < n and content[i] in "0123456789":
                    i += 1
                if i < n and content[i] == ".":
                    i += 1
                    while i < n and content[i] in "0123456789":
                        i += 1
                if i < n and content[i] in "eE":
                    i += 1
                    if i < n and content[i] in "+-":
                        i += 1
                    while i < n and content[i] in "0123456789":
                        i += 1

                tokens.append(JSONToken(content[start:i], JSONTokenType.NUMBER, start, i))
                continue

            # Boolean or null
            if content[i : i + 4] == "true":
                tokens.append(JSONToken("true", JSONTokenType.BOOLEAN, i, i + 4))
                i += 4
                continue
            if content[i : i + 5] == "false":
                tokens.append(JSONToken("false", JSONTokenType.BOOLEAN, i, i + 5))
                i += 5
                continue
            if content[i : i + 4] == "null":
                tokens.append(JSONToken("null", JSONTokenType.NULL, i, i + 4))
                i += 4
                continue

            # Unknown character - skip
            i += 1

        return tokens


def extract_json_schema(content: str) -> dict[str, Any] | list[Any]:
    """Extract the schema (keys only) from JSON content.

    Useful for understanding the structure without the values.

    Args:
        content: JSON content.

    Returns:
        Schema dictionary with keys and types (no values).

    Example:
        >>> extract_json_schema('{"name": "Alice", "age": 30}')
        {'name': 'string', 'age': 'number'}
    """

    def _extract(obj: Any) -> Any:
        if isinstance(obj, dict):
            return {k: _extract(v) for k, v in obj.items()}
        elif isinstance(obj, list):
            if obj:
                return [_extract(obj[0])]  # Schema of first item
            return []
        elif isinstance(obj, str):
            return "string"
        elif isinstance(obj, bool):
            return "boolean"
        elif isinstance(obj, int):
            return "integer"
        elif isinstance(obj, float):
            return "number"
        elif obj is None:
            return "null"
        else:
            return "unknown"

    try:
        parsed = json.loads(content)
        result = _extract(parsed)
        if isinstance(result, dict):
            return cast(dict[str, Any], result)
        elif isinstance(result, list):
            return cast(list[Any], result)
        return {}
    except (json.JSONDecodeError, ValueError):
        return {}
