"""
Markdown-aware chunking implementation.

Provides intelligent chunking for markdown content that respects
heading hierarchy and preserves document structure.
"""

import logging
import re
from typing import Any

from chunking.base import BaseChunker
from config import Settings
from models import Chunk, ChunkType, FileType

logger = logging.getLogger(__name__)

# Patterns for markdown elements
HEADING_PATTERN = re.compile(r"^(#{1,6})\s+(.+)$", re.MULTILINE)
CODE_BLOCK_PATTERN = re.compile(r"^```", re.MULTILINE)
HR_PATTERN = re.compile(r"^(-{3,}|_{3,}|\*{3,})$", re.MULTILINE)


class MarkdownChunker(BaseChunker):
    """
    Markdown-aware chunker that respects document structure.

    Features:
    - Respects heading hierarchy
    - Preserves heading context in chunks
    - Handles code blocks as units
    - Maintains list continuity where possible
    """

    def __init__(
        self,
        chunk_size: int,
        chunk_overlap: int,
        settings: Settings | None = None,
    ) -> None:
        """Initialize markdown chunker."""
        super().__init__(chunk_size, chunk_overlap, settings)

    @property
    def chunk_type(self) -> ChunkType:
        """Get chunk type."""
        return ChunkType.MARKDOWN

    def chunk(
        self,
        content: str,
        source_path: str | None = None,
        file_type: FileType | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> list[Chunk]:
        """
        Chunk markdown content.

        Splits on heading boundaries and preserves heading context.
        """
        if not content.strip():
            return []

        metadata = metadata or {}
        file_type = file_type or FileType.MARKDOWN

        # Split content into sections by headings
        sections = self._split_by_headings(content)

        if not sections:
            # No headings, chunk as plain text
            return self._chunk_text_block(content, source_path, file_type, metadata, [])

        chunks: list[Chunk] = []
        heading_stack: list[tuple[int, str]] = []  # (level, text)

        for section in sections:
            heading_level = section.get("level", 0)
            heading_text = section.get("heading", "")
            section_content = section.get("content", "")
            start_line = section.get("start_line", 1)
            end_line = section.get("end_line", 1)

            # Update heading stack
            if heading_level > 0:
                # Pop headings of equal or higher level
                while heading_stack and heading_stack[-1][0] >= heading_level:
                    heading_stack.pop()
                heading_stack.append((heading_level, heading_text))

            # Build heading context prefix
            heading_context = " > ".join(h[1] for h in heading_stack)

            section_chunks = self._chunk_section(
                content=section_content,
                heading_context=heading_context,
                heading_level=heading_level,
                heading_text=heading_text,
                start_line=start_line,
                end_line=end_line,
                source_path=source_path,
                file_type=file_type,
                metadata=metadata,
            )
            chunks.extend(section_chunks)

        return chunks

    def _split_by_headings(self, content: str) -> list[dict[str, Any]]:
        """Split content into sections by headings."""
        sections: list[dict[str, Any]] = []
        lines = content.split("\n")

        current_section: dict[str, Any] = {
            "level": 0,
            "heading": "",
            "content": "",
            "start_line": 1,
            "end_line": 1,
        }
        current_lines: list[str] = []
        in_code_block = False

        for i, line in enumerate(lines):
            # Track code blocks
            if line.strip().startswith("```"):
                in_code_block = not in_code_block
                current_lines.append(line)
                continue

            # Skip heading detection in code blocks
            if in_code_block:
                current_lines.append(line)
                continue

            # Check for heading
            heading_match = HEADING_PATTERN.match(line)
            if heading_match:
                # Save previous section
                if current_lines:
                    current_section["content"] = "\n".join(current_lines)
                    current_section["end_line"] = i
                    if current_section["content"].strip():
                        sections.append(current_section)

                # Start new section
                level = len(heading_match.group(1))
                heading_text = heading_match.group(2).strip()
                current_section = {
                    "level": level,
                    "heading": heading_text,
                    "content": "",
                    "start_line": i + 1,
                    "end_line": i + 1,
                }
                current_lines = [line]
            else:
                current_lines.append(line)

        # Save final section
        if current_lines:
            current_section["content"] = "\n".join(current_lines)
            current_section["end_line"] = len(lines)
            if current_section["content"].strip():
                sections.append(current_section)

        return sections

    def _chunk_section(
        self,
        content: str,
        heading_context: str,
        heading_level: int,
        heading_text: str,
        start_line: int,
        end_line: int,
        source_path: str | None,
        file_type: FileType,
        metadata: dict[str, Any],
    ) -> list[Chunk]:
        """Chunk a single section of markdown."""
        if not content.strip():
            return []

        token_count = self.count_tokens(content)

        # If section fits in one chunk, return as-is
        if token_count <= self.chunk_size:
            section_metadata = {
                **metadata,
                "heading_context": heading_context,
                "heading_level": heading_level,
                "heading_text": heading_text,
            }
            return [
                self._create_chunk(
                    content=content.strip(),
                    source_path=source_path,
                    start_line=start_line,
                    end_line=end_line,
                    file_type=file_type,
                    metadata=section_metadata,
                )
            ]

        # Need to split - try to split on paragraphs first
        return self._chunk_text_block(
            content,
            source_path,
            file_type,
            {
                **metadata,
                "heading_context": heading_context,
                "heading_level": heading_level,
                "heading_text": heading_text,
            },
            _heading_stack=[(heading_level, heading_text)] if heading_text else [],
            base_line=start_line,
        )

    def _chunk_text_block(
        self,
        content: str,
        source_path: str | None,
        file_type: FileType,
        metadata: dict[str, Any],
        _heading_stack: list[tuple[int, str]],
        base_line: int = 1,
    ) -> list[Chunk]:
        """Chunk a block of text by paragraphs."""
        # Split into paragraphs (separated by blank lines)
        paragraphs = self._split_into_paragraphs(content)

        if not paragraphs:
            return []

        chunks: list[Chunk] = []
        current_content: list[str] = []
        current_tokens = 0
        chunk_start_line = base_line

        for para_info in paragraphs:
            para_content = para_info["content"]
            para_tokens = para_info["tokens"]
            para_start = para_info["start_line"]

            # Handle very large paragraphs
            if para_tokens > self.chunk_size:
                # Flush current content
                if current_content:
                    chunk_text = "\n\n".join(current_content)
                    chunks.append(
                        self._create_chunk(
                            content=chunk_text.strip(),
                            source_path=source_path,
                            start_line=chunk_start_line,
                            end_line=base_line + para_start - 1,
                            file_type=file_type,
                            metadata=metadata,
                        )
                    )
                    current_content = []
                    current_tokens = 0

                # Split large paragraph by sentences/lines
                sub_chunks = self._split_large_paragraph(
                    para_content,
                    source_path,
                    file_type,
                    metadata,
                    base_line + para_start,
                )
                chunks.extend(sub_chunks)
                chunk_start_line = base_line + para_info["end_line"] + 1
                continue

            # Check if adding this paragraph exceeds limit
            if current_tokens + para_tokens > self.chunk_size and current_content:
                # Create chunk
                chunk_text = "\n\n".join(current_content)
                chunks.append(
                    self._create_chunk(
                        content=chunk_text.strip(),
                        source_path=source_path,
                        start_line=chunk_start_line,
                        end_line=base_line + para_start - 1,
                        file_type=file_type,
                        metadata=metadata,
                    )
                )

                # Overlap: include last paragraph if it fits
                if (
                    current_content
                    and self.count_tokens(current_content[-1]) <= self.chunk_overlap
                ):
                    current_content = [current_content[-1]]
                    current_tokens = self.count_tokens(current_content[-1])
                else:
                    current_content = []
                    current_tokens = 0

                chunk_start_line = base_line + para_start

            current_content.append(para_content)
            current_tokens += para_tokens

        # Final chunk
        if current_content:
            chunk_text = "\n\n".join(current_content)
            end_line_num = base_line + (paragraphs[-1]["end_line"] if paragraphs else 0)
            chunks.append(
                self._create_chunk(
                    content=chunk_text.strip(),
                    source_path=source_path,
                    start_line=chunk_start_line,
                    end_line=end_line_num,
                    file_type=file_type,
                    metadata=metadata,
                )
            )

        return chunks

    def _split_into_paragraphs(self, content: str) -> list[dict[str, Any]]:
        """Split content into paragraphs with metadata."""
        paragraphs: list[dict[str, Any]] = []
        lines = content.split("\n")

        current_para: list[str] = []
        para_start = 0
        in_code_block = False

        for i, line in enumerate(lines):
            # Track code blocks (keep them as single units)
            if line.strip().startswith("```"):
                if in_code_block:
                    # End of code block
                    current_para.append(line)
                    in_code_block = False
                else:
                    # Start of code block - save previous paragraph
                    if current_para and any(p.strip() for p in current_para):
                        para_content = "\n".join(current_para)
                        paragraphs.append(
                            {
                                "content": para_content,
                                "tokens": self.count_tokens(para_content),
                                "start_line": para_start,
                                "end_line": i - 1,
                            }
                        )
                    current_para = [line]
                    para_start = i
                    in_code_block = True
                continue

            if in_code_block:
                current_para.append(line)
                continue

            # Empty line indicates paragraph break
            if not line.strip():
                if current_para and any(p.strip() for p in current_para):
                    para_content = "\n".join(current_para)
                    paragraphs.append(
                        {
                            "content": para_content,
                            "tokens": self.count_tokens(para_content),
                            "start_line": para_start,
                            "end_line": i - 1,
                        }
                    )
                current_para = []
                para_start = i + 1
            else:
                if not current_para:
                    para_start = i
                current_para.append(line)

        # Final paragraph
        if current_para and any(p.strip() for p in current_para):
            para_content = "\n".join(current_para)
            paragraphs.append(
                {
                    "content": para_content,
                    "tokens": self.count_tokens(para_content),
                    "start_line": para_start,
                    "end_line": len(lines) - 1,
                }
            )

        return paragraphs

    def _split_large_paragraph(
        self,
        content: str,
        source_path: str | None,
        file_type: FileType,
        metadata: dict[str, Any],
        base_line: int,
    ) -> list[Chunk]:
        """Split a large paragraph into smaller chunks."""
        # Try splitting by sentences
        sentences = self._split_into_sentences(content)

        chunks: list[Chunk] = []
        current_content: list[str] = []
        current_tokens = 0

        for sentence in sentences:
            sentence_tokens = self.count_tokens(sentence)

            # If single sentence is too large, truncate
            if sentence_tokens > self.chunk_size:
                if current_content:
                    chunk_text = " ".join(current_content)
                    chunks.append(
                        self._create_chunk(
                            content=chunk_text.strip(),
                            source_path=source_path,
                            start_line=base_line,
                            end_line=base_line,
                            file_type=file_type,
                            metadata=metadata,
                        )
                    )
                    current_content = []
                    current_tokens = 0

                truncated = self.truncate_to_tokens(sentence, self.chunk_size)
                chunks.append(
                    self._create_chunk(
                        content=truncated.strip(),
                        source_path=source_path,
                        start_line=base_line,
                        end_line=base_line,
                        file_type=file_type,
                        metadata={**metadata, "truncated": True},
                    )
                )
                continue

            if current_tokens + sentence_tokens > self.chunk_size and current_content:
                chunk_text = " ".join(current_content)
                chunks.append(
                    self._create_chunk(
                        content=chunk_text.strip(),
                        source_path=source_path,
                        start_line=base_line,
                        end_line=base_line,
                        file_type=file_type,
                        metadata=metadata,
                    )
                )

                # Overlap with last sentence
                if (
                    current_content
                    and self.count_tokens(current_content[-1]) <= self.chunk_overlap
                ):
                    current_content = [current_content[-1]]
                    current_tokens = self.count_tokens(current_content[-1])
                else:
                    current_content = []
                    current_tokens = 0

            current_content.append(sentence)
            current_tokens += sentence_tokens

        # Final chunk
        if current_content:
            chunk_text = " ".join(current_content)
            chunks.append(
                self._create_chunk(
                    content=chunk_text.strip(),
                    source_path=source_path,
                    start_line=base_line,
                    end_line=base_line,
                    file_type=file_type,
                    metadata=metadata,
                )
            )

        return chunks

    def _split_into_sentences(self, text: str) -> list[str]:
        """Split text into sentences."""
        # Simple sentence splitting on common terminators
        # More sophisticated splitting could use nltk or spacy
        sentence_endings = re.compile(r"(?<=[.!?])\s+")
        sentences = sentence_endings.split(text)
        return [s.strip() for s in sentences if s.strip()]
