from __future__ import annotations

from typing import Any

from .utils import word_count, stable_id


def _split_words(text: str, max_words: int) -> list[str]:
    """Split very large single-message documents into bounded word chunks."""
    words = text.split()
    if len(words) <= max_words:
        return [text]
    return [" ".join(words[i:i + max_words]) for i in range(0, len(words), max_words)]


def chunk_conversation(conversation: dict[str, Any], max_words: int = 900) -> list[dict[str, Any]]:
    messages = conversation.get("messages") or []
    if not messages:
        return []

    chunks: list[dict[str, Any]] = []
    buffer: list[dict[str, Any]] = []
    buffer_words = 0
    start_idx = 1
    base_metadata = dict(conversation.get("metadata", {}) or {})
    base_metadata.update({
        "raw_path": conversation.get("raw_path"),
        "conversation_created_at": conversation.get("created_at"),
        "conversation_updated_at": conversation.get("updated_at"),
    })

    def add_chunk(text: str, start: int, end: int, extra_metadata: dict[str, Any] | None = None) -> None:
        if not text.strip():
            return
        metadata = dict(base_metadata)
        if extra_metadata:
            metadata.update(extra_metadata)
        chunk_id = stable_id(conversation.get("source", ""), conversation.get("conversation_id", ""), str(start), str(end), text[:200], str(metadata.get("part", "")))
        chunks.append({
            "chunk_id": chunk_id,
            "conversation_id": conversation.get("conversation_id"),
            "source": conversation.get("source"),
            "title": conversation.get("title"),
            "message_start": start,
            "message_end": end,
            "text": text,
            "created_at": conversation.get("created_at"),
            "updated_at": conversation.get("updated_at"),
            "metadata": metadata,
        })

    def flush(end_idx: int) -> None:
        nonlocal buffer, buffer_words, start_idx
        if not buffer:
            return
        text = "\n\n".join(f"{m.get('role', 'unknown')}: {m.get('content', '')}" for m in buffer).strip()
        add_chunk(text, start_idx, end_idx)
        buffer = []
        buffer_words = 0
        start_idx = end_idx + 1

    for idx, msg in enumerate(messages, start=1):
        content = str(msg.get("content") or "")
        msg_words = word_count(content)
        if msg_words > max_words:
            flush(idx - 1)
            role = msg.get("role", "unknown")
            pieces = _split_words(content, max_words)
            for part_idx, piece in enumerate(pieces, start=1):
                add_chunk(f"{role}: {piece}".strip(), idx, idx, {
                    "part": part_idx,
                    "parts_total": len(pieces),
                    "split_reason": "oversized_message",
                })
            start_idx = idx + 1
            continue
        if buffer and buffer_words + msg_words > max_words:
            flush(idx - 1)
        buffer.append(msg)
        buffer_words += msg_words
    flush(len(messages))
    return chunks