import io
import json
import re
from collections.abc import AsyncIterable
from dataclasses import dataclass
from typing import IO, Any, Callable

import ijson  # type: ignore
from inspect_ai._util.async_bytes_reader import AsyncBytesReader, adapt_to_reader
from inspect_ai.event import timeline_load
from pydantic import JsonValue

from ..types import (
    EventFilter,
    MessageFilter,
    Transcript,
    TranscriptContent,
    TranscriptInfo,
)
from ..util import filter_transcript
from .pool import resolve_pools
from .reducer import (
    ATTACHMENT_PREFIX,
    ATTACHMENTS_PREFIX,
    CALL_POOL_ITEM_PREFIX,
    EVENTS_ITEM_PREFIX,
    MESSAGE_POOL_ITEM_PREFIX,
    MESSAGES_ITEM_PREFIX,
    METADATA_PREFIX,
    SCORES_PREFIX,
    TARGET_PREFIX,
    TIMELINES_ITEM_PREFIX,
    ListProcessingConfig,
    ParseState,
    attachments_coroutine,
    call_pool_item_coroutine,
    event_item_coroutine,
    message_item_coroutine,
    message_pool_item_coroutine,
    metadata_coroutine,
    scores_coroutine,
    target_coroutine,
    timeline_item_coroutine,
)

# Pre-compiled regex patterns for performance
ATTACHMENT_PATTERN = re.compile(r"attachment://([a-f0-9]{32})")

# Section constants for prefix classification
_SECTION_OTHER = 0
_SECTION_MESSAGES = 1
_SECTION_EVENTS = 2
_SECTION_ATTACHMENTS = 3
_SECTION_METADATA = 4
_SECTION_TIMELINES = 5
_SECTION_TARGET = 6
_SECTION_SCORES = 7
_SECTION_MESSAGE_POOL = 8
_SECTION_CALL_POOL = 9

_MESSAGES_ITEM_PREFIX_LEN = len(MESSAGES_ITEM_PREFIX)
_EVENTS_ITEM_PREFIX_LEN = len(EVENTS_ITEM_PREFIX)
_ATTACHMENTS_PREFIX_LEN = len(ATTACHMENTS_PREFIX)
_METADATA_PREFIX_LEN = len(METADATA_PREFIX)
_TIMELINES_ITEM_PREFIX_LEN = len(TIMELINES_ITEM_PREFIX)
_SCORES_PREFIX_LEN = len(SCORES_PREFIX)
_TARGET_PREFIX_LEN = len(TARGET_PREFIX)
_MESSAGE_POOL_ITEM_PREFIX_LEN = len(MESSAGE_POOL_ITEM_PREFIX)
_CALL_POOL_ITEM_PREFIX_LEN = len(CALL_POOL_ITEM_PREFIX)
# "target" vs "timelines" — discriminate on 2nd char (derived from constant)
_TARGET_CHAR1 = TARGET_PREFIX[1]
_MIN_SECTION_PREFIX_LEN = min(
    _MESSAGES_ITEM_PREFIX_LEN,
    _EVENTS_ITEM_PREFIX_LEN,
    _ATTACHMENTS_PREFIX_LEN,
    _METADATA_PREFIX_LEN,
    _TIMELINES_ITEM_PREFIX_LEN,
    _SCORES_PREFIX_LEN,
    _TARGET_PREFIX_LEN,
    _MESSAGE_POOL_ITEM_PREFIX_LEN,
    _CALL_POOL_ITEM_PREFIX_LEN,
)


@dataclass(slots=True)
class RawTranscript:
    """Temporary structure for transcript data before validation."""

    id: str
    source_type: str | None
    source_id: str | None
    source_uri: str | None
    date: str | None
    task_set: str | None
    task_id: str | None
    task_repeat: int | None
    agent: str | None
    agent_args: dict[str, Any] | None
    model: str | None
    model_options: dict[str, Any] | None
    score: JsonValue
    success: bool | None
    message_count: int | None
    total_time: float | None
    total_tokens: int | None
    error: str | None
    limit: str | None
    metadata: dict[str, Any]
    messages: list[dict[str, Any]]
    events: list[dict[str, Any]]
    timelines: list[dict[str, Any]] | None = None


async def load_filtered_transcript(
    sample_bytes: IO[bytes] | AsyncIterable[bytes],
    t: TranscriptInfo,
    messages: MessageFilter,
    events: EventFilter,
    *,
    on_early_exit: Callable[[], None] | None = None,
) -> Transcript:
    """
    Transform and filter JSON sample data into a Transcript.

    Uses a two-phase approach:
    1. Stream parse and filter messages/events while collecting attachment references
    2. Resolve attachment references with actual values

    Falls back to non-streaming json5 parser if streaming fails (e.g., NaN/Inf values).

    Args:
        sample_bytes: Byte stream of JSON sample data
        t: TranscriptInfo representing the transcript to load
        messages: Filter for message roles (None=exclude all, "all"=include all,
            list=include matching)
        events: Filter for event types (None=exclude all, "all"=include all,
            list=include matching)
        on_early_exit: Test-only callback invoked immediately before the
            early-exit break

    Returns:
        Transcript object with filtered messages and events, resolved attachments.
        Metadata includes sample_metadata, target, and scores from the sample JSON.
        ``input`` is not unthinned: the sample JSON's input can contain attachment
        refs whose resolution requires parsing the attachments section — which follows
        events, defeating the early-exit optimization.
    """
    try:
        # Phase 1: Parse, filter, and collect attachment references
        async with adapt_to_reader(sample_bytes) as reader:
            transcript, attachment_refs = await _parse_and_filter(
                reader, t, messages, events, on_early_exit=on_early_exit
            )
        # Phase 2: Resolve attachment references
        return _resolve_attachments(transcript, attachment_refs)
    except ijson.JSONError:
        # Fallback to json5 for JSON5 features (NaN, Inf, etc.)
        return await _load_with_json5_fallback(sample_bytes, t, messages, events)


async def _load_with_json5_fallback(
    sample_bytes: IO[bytes] | AsyncIterable[bytes],
    t: TranscriptInfo,
    messages: MessageFilter,
    events: EventFilter,
) -> Transcript:
    """Fallback parser using json5 for JSON5 features (NaN, Inf, etc.)."""
    if hasattr(sample_bytes, "__aiter__"):
        io_source: IO[bytes] = io.BytesIO()
        async for chunk in sample_bytes:
            io_source.write(chunk)
    else:
        io_source = sample_bytes
    io_source.seek(0)
    data = json.load(io.TextIOWrapper(io_source, encoding="utf-8"))
    _resolve_pools_from_dict(data)

    return filter_transcript(
        _resolve_attachments(
            RawTranscript(
                id=t.transcript_id,
                source_type=t.source_type,
                source_id=t.source_id,
                source_uri=t.source_uri,
                date=t.date,
                task_set=t.task_set,
                task_id=t.task_id,
                task_repeat=t.task_repeat,
                agent=t.agent,
                agent_args=t.agent_args,
                model=t.model,
                model_options=t.model_options,
                score=t.score,
                success=t.success,
                message_count=t.message_count,
                total_time=t.total_time,
                total_tokens=t.total_tokens,
                error=t.error,
                limit=t.limit,
                metadata=_merge_unthinned_from_dict(t.metadata, data),
                messages=data.get("messages", []),
                events=data.get("events", []),
                timelines=data.get("timelines"),
            ),
            data.get("attachments", {}),
        ),
        TranscriptContent(messages, events),
    )


def _merge_unthinned(base: dict[str, Any], state: ParseState) -> dict[str, Any]:
    """Merge unthinned fields from stream parse into transcript metadata."""
    overrides: dict[str, Any] = {}
    if state.metadata:
        overrides["sample_metadata"] = state.metadata
    if state.target is not None:
        overrides["target"] = state.target
    if state.scores:
        overrides["scores"] = state.scores
    return base.copy() | overrides if overrides else base


def _resolve_pools_from_dict(data: dict[str, Any]) -> None:
    """Resolve v3 pools in a fully-parsed sample dict (mutates in-place).

    No-op when pools are absent or empty (v2 files).
    """
    from .pool import _resolve_events_pools

    _resolve_events_pools(
        data.get("events", []),
        data.get("message_pool", []),
        data.get("call_pool", []),
    )


def _merge_unthinned_from_dict(
    base: dict[str, Any], data: dict[str, Any]
) -> dict[str, Any]:
    """Merge unthinned fields from fully-parsed dict into transcript metadata."""
    overrides: dict[str, Any] = {}
    if data.get("metadata"):
        overrides["sample_metadata"] = data["metadata"]
    if "target" in data:
        overrides["target"] = data["target"]
    if data.get("scores"):
        overrides["scores"] = data["scores"]
    return base.copy() | overrides if overrides else base


async def _parse_and_filter(
    sample_json: AsyncBytesReader,
    t: TranscriptInfo,
    messages_filter: MessageFilter,
    events_filter: EventFilter,
    *,
    on_early_exit: Callable[[], None] | None = None,
) -> tuple[RawTranscript, dict[str, str]]:
    """
    Phase 1: Single-pass stream parse, filter, and collect attachment references.

    Returns:
        Tuple of (partial transcript, attachment references dict)
    """
    # Create processing configurations
    messages_config = (
        ListProcessingConfig(
            array_item_prefix="messages.item",
            filter_field="role",
            filter_list=messages_filter,  # type:ignore
        )
        if messages_filter is not None
        else None
    )

    events_config = (
        ListProcessingConfig(
            array_item_prefix="events.item",
            filter_field="event",
            filter_list=events_filter,  # type:ignore
        )
        if events_filter is not None
        else None
    )

    state = ParseState()

    # Initialize coroutine processors
    messages_coro = (
        message_item_coroutine(state, messages_config) if messages_config else None
    )
    events_coro = event_item_coroutine(state, events_config) if events_config else None
    timelines_coro = timeline_item_coroutine(state)
    attachments_coro = attachments_coroutine(state)
    metadata_coro = metadata_coroutine(state)
    target_coro = target_coroutine(state)
    scores_coro = scores_coroutine(state)
    message_pool_coro = message_pool_item_coroutine(state) if events_coro else None
    call_pool_coro = call_pool_item_coroutine(state) if events_coro else None

    last_prefix = ""
    current_section = _SECTION_OTHER

    async for prefix, event, value in ijson.parse_async(sample_json, use_float=True):
        # Early exit: skip events/attachments when they aren't needed.
        # JSON field order is: ...target, messages, output, scores, metadata,
        # store, events, attachments — so by the time we see "events" start_array,
        # metadata and scores have already been parsed.
        if (
            events_coro is None
            and prefix == "events"
            and event == "start_array"
            and not state.attachment_refs
        ):
            if on_early_exit is not None:
                on_early_exit()
            break

        # Inline prefix classification for performance (56M+ calls in hot path).
        # WARNING: every operation here can run millions of times per parse.
        # Avoid string slicing, startswith, or any allocation in common paths.
        # Profile before changing.
        if prefix != last_prefix:
            last_prefix = prefix
            p_len = len(prefix)
            if p_len == 0 or prefix[0] not in ("m", "e", "a", "t", "s", "c"):
                current_section = _SECTION_OTHER
            elif p_len < _MIN_SECTION_PREFIX_LEN:
                # Short prefixes: "scores" (6), "target" (6)
                if prefix == "scores":
                    current_section = _SECTION_SCORES
                elif prefix == "target":
                    current_section = _SECTION_TARGET
                else:
                    current_section = _SECTION_OTHER
            elif prefix[0] == "m":
                # Both "messages" and "metadata" start with "me", check 3rd char
                # (safe because we already checked p_len >= _MIN_SECTION_PREFIX_LEN)
                if (
                    p_len >= _MESSAGES_ITEM_PREFIX_LEN
                    and prefix[2] == "s"
                    and prefix[:_MESSAGES_ITEM_PREFIX_LEN] == MESSAGES_ITEM_PREFIX
                ):
                    current_section = _SECTION_MESSAGES
                elif prefix[2] == "t" and (
                    prefix == "metadata" or prefix.startswith(METADATA_PREFIX)
                ):
                    current_section = _SECTION_METADATA
                elif (
                    p_len >= _MESSAGE_POOL_ITEM_PREFIX_LEN
                    and prefix[:_MESSAGE_POOL_ITEM_PREFIX_LEN]
                    == MESSAGE_POOL_ITEM_PREFIX
                ):
                    current_section = _SECTION_MESSAGE_POOL
                else:
                    current_section = _SECTION_OTHER
            elif (
                prefix[0] == "e"
                and p_len >= _EVENTS_ITEM_PREFIX_LEN
                and prefix[:_EVENTS_ITEM_PREFIX_LEN] == EVENTS_ITEM_PREFIX
            ):
                current_section = _SECTION_EVENTS
            elif (
                prefix[0] == "a"
                and p_len >= _ATTACHMENTS_PREFIX_LEN
                and prefix[:_ATTACHMENTS_PREFIX_LEN] == ATTACHMENTS_PREFIX
            ):
                current_section = _SECTION_ATTACHMENTS
            elif (
                prefix[0] == "c"
                and p_len >= _CALL_POOL_ITEM_PREFIX_LEN
                and prefix[:_CALL_POOL_ITEM_PREFIX_LEN] == CALL_POOL_ITEM_PREFIX
            ):
                current_section = _SECTION_CALL_POOL
            elif prefix[0] == "t":
                if prefix[1] == _TARGET_CHAR1:
                    current_section = _SECTION_TARGET
                elif (
                    p_len >= _TIMELINES_ITEM_PREFIX_LEN
                    and prefix[:_TIMELINES_ITEM_PREFIX_LEN] == TIMELINES_ITEM_PREFIX
                ):
                    current_section = _SECTION_TIMELINES
                else:
                    current_section = _SECTION_OTHER
            elif prefix[0] == "s" and prefix[:_SCORES_PREFIX_LEN] == SCORES_PREFIX:
                current_section = _SECTION_SCORES
            else:
                current_section = _SECTION_OTHER

        # Dispatch to coroutines (optimized to avoid redundant None checks)
        if current_section == _SECTION_MESSAGES and messages_coro:
            messages_coro.send((prefix, event, value))
        elif current_section == _SECTION_EVENTS and events_coro:
            events_coro.send((prefix, event, value))
        elif current_section == _SECTION_ATTACHMENTS:
            attachments_coro.send((prefix, event, value))
        elif current_section == _SECTION_METADATA:
            metadata_coro.send((prefix, event, value))
        elif current_section == _SECTION_TARGET and target_coro is not None:
            try:
                target_coro.send((prefix, event, value))
            except StopIteration:
                target_coro = None
        elif current_section == _SECTION_TIMELINES:
            timelines_coro.send((prefix, event, value))
        elif current_section == _SECTION_SCORES:
            scores_coro.send((prefix, event, value))
        elif current_section == _SECTION_MESSAGE_POOL and message_pool_coro:
            message_pool_coro.send((prefix, event, value))
        elif current_section == _SECTION_CALL_POOL and call_pool_coro:
            call_pool_coro.send((prefix, event, value))

    # Resolve v3 pool references before returning
    resolve_pools(state)

    return (
        RawTranscript(
            id=t.transcript_id,
            source_type=t.source_type,
            source_id=t.source_id,
            source_uri=t.source_uri,
            date=t.date,
            task_set=t.task_set,
            task_id=t.task_id,
            task_repeat=t.task_repeat,
            agent=t.agent,
            agent_args=t.agent_args,
            model=t.model,
            model_options=t.model_options,
            score=t.score,
            success=t.success,
            message_count=t.message_count,
            total_time=t.total_time,
            total_tokens=t.total_tokens,
            error=t.error,
            limit=t.limit,
            metadata=_merge_unthinned(t.metadata, state),
            messages=state.messages,
            events=state.events,
            timelines=state.timelines if state.timelines else None,
        ),
        state.attachments,
    )


def _resolve_attachments(
    transcript: RawTranscript, attachments: dict[str, str]
) -> Transcript:
    """
    Phase 2: Replace attachment references with actual values.

    Args:
        transcript: Transcript with attachment references
        attachments: Dict mapping attachment IDs to their values

    Returns:
        Transcript with resolved attachment references
    """

    def resolve_string(text: str) -> str:
        """Replace attachment references in a string."""
        # Fast path: skip regex if no attachment prefix found
        if ATTACHMENT_PREFIX not in text:
            return text

        def replace_ref(match: re.Match[str]) -> str:
            attachment_id = match.group(1)
            return attachments.get(
                attachment_id, match.group(0)
            )  # Return original if not found

        return ATTACHMENT_PATTERN.sub(replace_ref, text)

    # Resolve references in messages (already raw dicts, no need to model_dump)
    resolved_messages = []
    for message_dict in transcript.messages:
        resolved_dict = _resolve_dict_attachments(message_dict, resolve_string)
        resolved_messages.append(resolved_dict)

    # Resolve references in events (already raw dicts, no need to model_dump)
    resolved_events = []
    for event_dict in transcript.events:
        resolved_dict = _resolve_dict_attachments(event_dict, resolve_string)
        resolved_events.append(resolved_dict)

    # Create new transcript with resolved data
    # Use model_validate to validate messages/events into proper types,
    # but pass metadata separately via __pydantic_private__ to preserve LazyJSONDict
    transcript_data: dict[str, Any] = {
        "transcript_id": transcript.id,
        "source_type": transcript.source_type,
        "source_id": transcript.source_id,
        "source_uri": transcript.source_uri,
        "date": transcript.date,
        "task_set": transcript.task_set,
        "task_id": transcript.task_id,
        "task_repeat": transcript.task_repeat,
        "agent": transcript.agent,
        "agent_args": transcript.agent_args,
        "model": transcript.model,
        "model_options": transcript.model_options,
        "score": transcript.score,
        "success": transcript.success,
        "message_count": transcript.message_count,
        "total_time": transcript.total_time,
        "total_tokens": transcript.total_tokens,
        "error": transcript.error,
        "limit": transcript.limit,
        "metadata": {},  # Placeholder to avoid validation
        "messages": resolved_messages,
        "events": resolved_events,
    }

    validated = Transcript.model_validate(transcript_data)

    # Resolve timelines with event UUID context (events must be validated first)
    if transcript.timelines:
        resolved_timelines = [
            timeline_load(tl_dict, validated.events) for tl_dict in transcript.timelines
        ]
        validated.timelines = resolved_timelines

    # Directly assign metadata to preserve LazyJSONDict
    validated.metadata = transcript.metadata
    return validated


def _resolve_dict_attachments(obj: Any, resolve_func: Callable[[str], str]) -> Any:
    if isinstance(obj, str):
        return resolve_func(obj)
    if isinstance(obj, dict):
        return {k: _resolve_dict_attachments(v, resolve_func) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_resolve_dict_attachments(item, resolve_func) for item in obj]

    return obj
