from typing import List, Optional

from phi.document.chunking.strategy import ChunkingStrategy
from phi.document.base import Document
from phi.model.openai import OpenAIChat
from phi.model.base import Model
from phi.model.message import Message


class AgenticChunking(ChunkingStrategy):
    """Chunking strategy that uses an LLM to determine natural breakpoints in the text"""

    def __init__(self, model: Optional[Model] = None, max_chunk_size: int = 5000):
        self.model = model or OpenAIChat()
        self.max_chunk_size = max_chunk_size

    def chunk(self, document: Document) -> List[Document]:
        """Split text into chunks using LLM to determine natural breakpoints based on context"""
        if len(document.content) <= self.max_chunk_size:
            return [document]

        chunks: List[Document] = []
        remaining_text = self.clean_text(document.content)
        chunk_meta_data = document.meta_data
        chunk_number = 1

        while remaining_text:
            # Ask model to find a good breakpoint within max_chunk_size
            prompt = f"""Analyze this text and determine a natural breakpoint within the first {self.max_chunk_size} characters. 
            Consider semantic completeness, paragraph boundaries, and topic transitions.
            Return only the character position number of where to break the text:
            
            {remaining_text[: self.max_chunk_size]}"""

            try:
                response = self.model.response([Message(role="user", content=prompt)])
                if response and response.content:
                    break_point = min(int(response.content.strip()), self.max_chunk_size)
                else:
                    break_point = self.max_chunk_size
            except Exception:
                # Fallback to max size if model fails
                break_point = self.max_chunk_size

            # Extract chunk and update remaining text
            chunk = remaining_text[:break_point].strip()
            meta_data = chunk_meta_data.copy()
            meta_data["chunk"] = chunk_number
            chunk_id = None
            if document.id:
                chunk_id = f"{document.id}_{chunk_number}"
            elif document.name:
                chunk_id = f"{document.name}_{chunk_number}"
            meta_data["chunk_size"] = len(chunk)
            chunks.append(
                Document(
                    id=chunk_id,
                    name=document.name,
                    meta_data=meta_data,
                    content=chunk,
                )
            )
            chunk_number += 1

            remaining_text = remaining_text[break_point:].strip()

            if not remaining_text:
                break

        return chunks