"""Transform Cursor state.vscdb (cursorDiskKV) sessions to unified format."""

from __future__ import annotations

import base64
import hashlib
import json
from datetime import datetime, timezone
from typing import Any, Literal

from qc_trace.schemas.unified import (
    NormalizedMessage,
    TokenUsage,
    ToolCall,
    ToolResult,
)
from qc_trace.utils.vscdb import VscdbSession


def transform_cursor_vscdb(session: VscdbSession) -> list[NormalizedMessage]:
    """Transform a VscdbSession into normalized messages.

    Dispatches to mode-specific transform based on composerData structure.
    """
    mode = _detect_mode(session.composer_data)
    if mode == "inline":
        return _transform_inline(session)
    elif mode == "hashchain":
        return _transform_hashchain(session)
    else:
        return _transform_headers_only(session)


def _detect_mode(cd: dict) -> Literal["inline", "hashchain", "headers_only"]:
    """Detect which storage mode the composerData uses.

    - inline: has conversation[] with text content in bubbles
    - hashchain: has conversationState (base64-encoded protobuf with hashes pointing to agentKv blobs)
    - headers_only: has fullConversationHeadersOnly[] (bubble IDs only, no inline content)
    """
    if cd.get("conversation") and isinstance(cd["conversation"], list) and len(cd["conversation"]) > 0:
        # Check if bubbles actually have text content (not just headers)
        for bubble in cd["conversation"]:
            if isinstance(bubble, dict) and bubble.get("text"):
                return "inline"

    if cd.get("conversationState") and isinstance(cd["conversationState"], str):
        return "hashchain"

    if cd.get("fullConversationHeadersOnly") and isinstance(cd["fullConversationHeadersOnly"], list):
        return "headers_only"

    # Fallback: if conversation exists but no text, treat as headers_only
    if cd.get("conversation"):
        return "headers_only"

    return "headers_only"


def _transform_inline(session: VscdbSession) -> list[NormalizedMessage]:
    """Transform inline mode: conversation[] has full text content.

    Timestamps: timingInfo.clientStartTime for assistant bubbles,
    composerData.createdAt for user bubbles and fallback.
    Tool calls identified by isCapabilityIteration=True.
    Tokens from bubbleId lookup.
    """
    messages: list[NormalizedMessage] = []
    cd = session.composer_data
    conversation = cd.get("conversation", [])
    fallback_ts = _ms_to_iso(cd.get("createdAt", 0))
    model = _extract_model(cd)

    for i, bubble in enumerate(conversation):
        if not isinstance(bubble, dict):
            continue

        bubble_type = bubble.get("type", 0)
        bubble_id = bubble.get("bubbleId", "")
        text = bubble.get("text", "")
        timing = bubble.get("timingInfo") or {}
        is_capability = bubble.get("isCapabilityIteration", False)
        capability_type = bubble.get("capabilityType", "")

        # Determine timestamp
        client_start = timing.get("clientStartTime")
        ts = _ms_to_iso(client_start) if client_start else fallback_ts

        # Get token usage from bubble entry
        tokens = _bubble_tokens(session, session.composer_id, bubble_id)

        # Extract thinking
        thinking = None
        thinking_blocks = bubble.get("allThinkingBlocks", [])
        if thinking_blocks:
            parts = []
            for tb in thinking_blocks:
                if isinstance(tb, dict) and tb.get("thinking"):
                    parts.append(tb["thinking"])
            if parts:
                thinking = "\n\n".join(parts)

        if is_capability:
            # Tool call bubble
            msg_id = f"{session.composer_id}-{i}"
            tool_name = capability_type or "unknown"
            messages.append(NormalizedMessage(
                id=msg_id,
                session_id=session.composer_id,
                source="cursor_vscdb",
                source_schema_version=2,
                msg_type="tool_call",
                timestamp=ts,
                content=text if text else None,
                tokens=tokens,
                tool_call=ToolCall(
                    id=msg_id,
                    name=tool_name,
                    input={},
                ),
                model=model,
                raw_file_path=session.db_path,
            ))
        elif bubble_type == 1:
            # User message
            messages.append(NormalizedMessage(
                id=f"{session.composer_id}-{i}",
                session_id=session.composer_id,
                source="cursor_vscdb",
                source_schema_version=2,
                msg_type="user",
                timestamp=ts if not client_start else fallback_ts,
                content=text if text else None,
                tokens=tokens,
                model=model,
                raw_file_path=session.db_path,
            ))
        elif bubble_type == 2:
            # Assistant message
            messages.append(NormalizedMessage(
                id=f"{session.composer_id}-{i}",
                session_id=session.composer_id,
                source="cursor_vscdb",
                source_schema_version=2,
                msg_type="assistant",
                timestamp=ts,
                content=text if text else None,
                tokens=tokens,
                thinking=thinking,
                model=model,
                raw_file_path=session.db_path,
            ))
        else:
            # Unknown type, emit as assistant
            if text:
                messages.append(NormalizedMessage(
                    id=f"{session.composer_id}-{i}",
                    session_id=session.composer_id,
                    source="cursor_vscdb",
                    source_schema_version=2,
                    msg_type="assistant",
                    timestamp=ts,
                    content=text,
                    tokens=tokens,
                    model=model,
                    raw_file_path=session.db_path,
                ))

    return messages


def _transform_hashchain(session: VscdbSession) -> list[NormalizedMessage]:
    """Transform hashchain mode: conversationState → base64 → protobuf → SHA-256 → agentKv blobs.

    Timestamps: composerData.createdAt (session-level only).
    Tool results from agentKv role:tool entries with real status.
    """
    messages: list[NormalizedMessage] = []
    cd = session.composer_data
    fallback_ts = _ms_to_iso(cd.get("createdAt", 0))
    model = _extract_model(cd)

    conversation_state = cd.get("conversationState", "")
    hashes = _extract_hashes(conversation_state)

    for i, h in enumerate(hashes):
        kv_key = f"agentKv:blob:{h}"
        raw = session.agent_kv_entries.get(kv_key)
        if raw is None:
            continue

        parsed = _parse_agent_kv_json(raw)
        if parsed is None:
            continue

        role = parsed.get("role", "")
        msg_id = f"{session.composer_id}-hc-{i}"

        if role == "user":
            content_parts = parsed.get("content", [])
            text = _extract_text_from_content(content_parts)
            messages.append(NormalizedMessage(
                id=msg_id,
                session_id=session.composer_id,
                source="cursor_vscdb",
                source_schema_version=2,
                msg_type="user",
                timestamp=fallback_ts,
                content=text,
                model=model,
                raw_file_path=session.db_path,
            ))
        elif role == "assistant":
            content_parts = parsed.get("content", [])
            text = _extract_text_from_content(content_parts)
            tool_calls = _extract_tool_calls_from_content(content_parts, msg_id)

            if text:
                messages.append(NormalizedMessage(
                    id=msg_id,
                    session_id=session.composer_id,
                    source="cursor_vscdb",
                    source_schema_version=2,
                    msg_type="assistant",
                    timestamp=fallback_ts,
                    content=text,
                    model=model,
                    raw_file_path=session.db_path,
                ))

            for j, tc in enumerate(tool_calls):
                tc_id = f"{msg_id}-tc-{j}"
                messages.append(NormalizedMessage(
                    id=tc_id,
                    session_id=session.composer_id,
                    source="cursor_vscdb",
                    source_schema_version=2,
                    msg_type="tool_call",
                    timestamp=fallback_ts,
                    tool_call=tc,
                    model=model,
                    raw_file_path=session.db_path,
                ))
        elif role == "tool":
            tool_call_id = parsed.get("id") or parsed.get("tool_call_id") or msg_id
            content_parts = parsed.get("content", [])
            output = _extract_text_from_content(content_parts)
            # Determine status from content
            status: Literal["success", "failure"] = "success"
            if isinstance(content_parts, list):
                for part in content_parts:
                    if isinstance(part, dict) and part.get("is_error"):
                        status = "failure"
                        break
            messages.append(NormalizedMessage(
                id=msg_id,
                session_id=session.composer_id,
                source="cursor_vscdb",
                source_schema_version=2,
                msg_type="tool_result",
                timestamp=fallback_ts,
                tool_result=ToolResult(
                    call_id=tool_call_id,
                    output=output or "",
                    status=status,
                ),
                model=model,
                raw_file_path=session.db_path,
            ))

    return messages


def _transform_headers_only(session: VscdbSession) -> list[NormalizedMessage]:
    """Transform headers-only mode: fullConversationHeadersOnly with bubble ID lookups.

    Tokens from bubbleId entries. Minimal content (none from headers).
    Timestamps from bubbleId.createdAt (v3) or composerData.createdAt.
    """
    messages: list[NormalizedMessage] = []
    cd = session.composer_data
    fallback_ts = _ms_to_iso(cd.get("createdAt", 0))
    model = _extract_model(cd)

    headers = cd.get("fullConversationHeadersOnly") or cd.get("conversation") or []

    for i, bubble in enumerate(headers):
        if not isinstance(bubble, dict):
            continue

        bubble_type = bubble.get("type", 0)
        bubble_id = bubble.get("bubbleId", "")
        is_capability = bubble.get("isCapabilityIteration", False)

        # Look up bubble entry for tokens and timestamp
        tokens = _bubble_tokens(session, session.composer_id, bubble_id)
        bubble_entry = session.bubble_entries.get(
            f"bubbleId:{session.composer_id}:{bubble_id}"
        )

        ts = fallback_ts
        if bubble_entry and bubble_entry.get("createdAt"):
            ts = _ms_to_iso(bubble_entry["createdAt"])

        msg_id = f"{session.composer_id}-ho-{i}"

        if is_capability:
            capability_type = bubble.get("capabilityType", "unknown")
            messages.append(NormalizedMessage(
                id=msg_id,
                session_id=session.composer_id,
                source="cursor_vscdb",
                source_schema_version=2,
                msg_type="tool_call",
                timestamp=ts,
                tokens=tokens,
                tool_call=ToolCall(id=msg_id, name=capability_type, input={}),
                model=model,
                raw_file_path=session.db_path,
            ))
        elif bubble_type == 1:
            messages.append(NormalizedMessage(
                id=msg_id,
                session_id=session.composer_id,
                source="cursor_vscdb",
                source_schema_version=2,
                msg_type="user",
                timestamp=ts,
                tokens=tokens,
                model=model,
                raw_file_path=session.db_path,
            ))
        elif bubble_type == 2:
            messages.append(NormalizedMessage(
                id=msg_id,
                session_id=session.composer_id,
                source="cursor_vscdb",
                source_schema_version=2,
                msg_type="assistant",
                timestamp=ts,
                tokens=tokens,
                model=model,
                raw_file_path=session.db_path,
            ))

    return messages


# --- Helpers ---


def _extract_model(cd: dict) -> str | None:
    """Extract model name from composerData.modelConfig."""
    mc = cd.get("modelConfig")
    if isinstance(mc, dict):
        return mc.get("modelName")
    return None


def _ms_to_iso(ms: int | str) -> str:
    """Convert millisecond epoch to ISO 8601 string. Returns empty string for 0/None."""
    if not ms:
        return ""
    try:
        ms = int(ms)
        dt = datetime.fromtimestamp(ms / 1000, tz=timezone.utc)
        return dt.isoformat()
    except (OSError, ValueError, OverflowError):
        return ""


def _bubble_tokens(session: VscdbSession, composer_id: str, bubble_id: str) -> TokenUsage:
    """Look up token counts from the bubbleId entry."""
    if not bubble_id:
        return TokenUsage()
    key = f"bubbleId:{composer_id}:{bubble_id}"
    entry = session.bubble_entries.get(key)
    if not entry:
        return TokenUsage()
    tc = entry.get("tokenCount")
    if not isinstance(tc, dict):
        return TokenUsage()
    return TokenUsage(
        input=tc.get("inputTokens", 0),
        output=tc.get("outputTokens", 0),
    )


def _extract_hashes(conversation_state: str) -> list[str]:
    """Extract SHA-256 hashes from a base64-encoded protobuf conversation state.

    Format: optional leading '~', then base64 data.
    Walk the protobuf wire format looking for length-delimited fields of exactly 32 bytes.
    """
    if not conversation_state:
        return []

    # Strip leading tilde
    data_str = conversation_state.lstrip("~")
    if not data_str:
        return []

    # Add base64 padding
    padding = 4 - (len(data_str) % 4)
    if padding < 4:
        data_str += "=" * padding

    try:
        data = base64.b64decode(data_str)
    except Exception:
        return []

    return _walk_protobuf_for_hashes(data)


def _walk_protobuf_for_hashes(data: bytes) -> list[str]:
    """Walk protobuf wire format and extract 32-byte length-delimited fields as hex hashes."""
    hashes: list[str] = []
    pos = 0
    length = len(data)

    while pos < length:
        try:
            tag, pos = _decode_varint(data, pos)
        except (IndexError, ValueError):
            break

        wire_type = tag & 0x07

        if wire_type == 0:  # varint
            try:
                _, pos = _decode_varint(data, pos)
            except (IndexError, ValueError):
                break
        elif wire_type == 2:  # length-delimited
            try:
                field_len, pos = _decode_varint(data, pos)
            except (IndexError, ValueError):
                break
            if pos + field_len > length:
                break
            if field_len == 32:
                hash_bytes = data[pos : pos + 32]
                hashes.append(hash_bytes.hex())
            pos += field_len
        elif wire_type == 5:  # 32-bit fixed
            pos += 4
        elif wire_type == 1:  # 64-bit fixed
            pos += 8
        else:
            break  # unknown wire type

    return hashes


def _decode_varint(data: bytes, pos: int) -> tuple[int, int]:
    """Decode a protobuf varint starting at pos. Returns (value, new_pos)."""
    result = 0
    shift = 0
    while True:
        if pos >= len(data):
            raise IndexError("Varint extends beyond data")
        byte = data[pos]
        pos += 1
        result |= (byte & 0x7F) << shift
        if (byte & 0x80) == 0:
            break
        shift += 7
        if shift >= 64:
            raise ValueError("Varint too long")
    return result, pos


def _parse_agent_kv_json(raw: bytes) -> dict | None:
    """Decompress (if needed) and parse agent KV blob as JSON dict."""
    from qc_trace.utils.vscdb import _decompress
    try:
        decompressed = _decompress(raw)
        text = decompressed.decode("utf-8", errors="replace")
        parsed = json.loads(text)
        if isinstance(parsed, dict):
            return parsed
    except (json.JSONDecodeError, UnicodeDecodeError, ValueError):
        pass
    return None


def _extract_text_from_content(content: Any) -> str | None:
    """Extract text from a content array (list of content blocks) or string."""
    if isinstance(content, str):
        return content if content else None
    if not isinstance(content, list):
        return None
    parts = []
    for block in content:
        if isinstance(block, str):
            parts.append(block)
        elif isinstance(block, dict):
            if block.get("type") == "text" and block.get("text"):
                parts.append(block["text"])
    return "\n".join(parts) if parts else None


def _extract_tool_calls_from_content(content: Any, parent_id: str) -> list[ToolCall]:
    """Extract tool_use blocks from a content array."""
    if not isinstance(content, list):
        return []
    calls = []
    for block in content:
        if isinstance(block, dict) and block.get("type") == "tool_use":
            calls.append(ToolCall(
                id=block.get("id", parent_id),
                name=block.get("name", "unknown"),
                input=block.get("input", {}),
            ))
    return calls
