"""Ontology selection agent for OntoCast.

This module provides functionality for selecting appropriate ontologies based on
the content of source text segments, ensuring that the chosen ontology best matches the
domain and requirements of the text.
"""

import logging

from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import PromptTemplate

from ontocast.agent.common import call_llm_with_retry
from ontocast.onto.enum import Status
from ontocast.onto.model import create_ontology_selector_report_model
from ontocast.onto.null import NULL_ONTOLOGY
from ontocast.onto.state import AgentState
from ontocast.prompt.select_ontology import template_prompt
from ontocast.tool import OntologyManager
from ontocast.toolbox import ToolBox

logger = logging.getLogger(__name__)


def _create_document_excerpt(state: AgentState, max_length: int = 3000) -> str:
    """Create a representative excerpt from the document for ontology selection.

    This function samples text from multiple content units to provide a better
    representation of the document content than just the first unit.

    Args:
        state: The current agent state.
        max_length: Maximum total length of the excerpt.

    Returns:
        str: A representative excerpt from the document.
    """
    excerpt_parts = []
    total_length = 0
    chunk_length = max_length // 3  # Aim for ~3 source chunks, ~1000 chars each

    # Strategy: Sample from first, middle, and last units if available
    if state.content_units:
        num_chunks = len(state.content_units)
        indices_to_sample = []

        if num_chunks == 1:
            indices_to_sample = [0]
        elif num_chunks == 2:
            indices_to_sample = [0, 1]
        else:
            # Sample first, middle, and last
            indices_to_sample = [0, num_chunks // 2, num_chunks - 1]

        for idx in indices_to_sample:
            if idx < num_chunks and total_length < max_length:
                chunk_text = state.content_units[idx].text
                # Take a portion of this source chunk
                remaining = max_length - total_length
                sample_length = min(chunk_length, remaining, len(chunk_text))

                if sample_length > 0:
                    if sample_length < len(chunk_text):
                        excerpt_parts.append(chunk_text[:sample_length] + " ...")
                    else:
                        excerpt_parts.append(chunk_text)
                    total_length += sample_length

        if excerpt_parts:
            return "\n\n[...]\n\n".join(excerpt_parts)

    # Fallback: Use input_text if available
    if state.input_text:
        if len(state.input_text) <= max_length:
            return state.input_text
        return state.input_text[:max_length] + " ..."

    # Last resort: Use current content unit
    if state.current_content_unit and state.current_content_unit.text:
        chunk_text = state.current_content_unit.text
        if len(chunk_text) <= max_length:
            return chunk_text
        return chunk_text[:max_length] + " ..."

    return ""


async def select_ontology(state: AgentState, tools: ToolBox) -> AgentState:
    """Select an appropriate ontology for the document.

    This function analyzes the document and selects the most appropriate
    ontology based on its content and requirements using a numbered list selection.
    If an ontology is already selected, it skips selection to ensure one ontology
    per document.

    Args:
        state: The current agent state containing the document to process.
        tools: The toolbox instance providing utility functions.

    Returns:
        AgentState: Updated state with selected ontology.
    """
    # Skip if ontology already selected (for subsequent chunks in the loop)
    if not state.current_ontology.is_null():
        logger.debug(
            f"Ontology already selected: {state.current_ontology.ontology_id}, "
            "skipping selection to maintain one ontology per document"
        )
        state.status = Status.SUCCESS
        return state

    progress_info = state.get_content_unit_progress_string()
    logger.info(f"Selecting ontology for document ({progress_info})")
    llm_tool = tools.llm
    om_tool: OntologyManager = tools.ontology_manager

    if om_tool.has_ontologies:
        ontologies = om_tool.ontologies
        num_ontologies = len(ontologies)

        # Create numbered list of ontologies
        ontologies_list_lines = []
        for i, ontology in enumerate(ontologies, start=1):
            ontologies_list_lines.append(f"{i}. {ontology.describe()}")

        ontologies_list = "\n\n".join(ontologies_list_lines)

        logger.info(f"Presenting {num_ontologies} ontologies for selection")

        # Create a better document excerpt using multiple chunks
        excerpt = _create_document_excerpt(state, max_length=3000)

        # Create dynamic model with correct constraint
        ontology_selector_report_model = create_ontology_selector_report_model(
            num_ontologies
        )
        parser = PydanticOutputParser(pydantic_object=ontology_selector_report_model)

        prompt = PromptTemplate(
            template=template_prompt,
            input_variables=[
                "excerpt",
                "ontologies_list",
                "num_ontologies",
                "format_instructions",
            ],
        )

        selector = await call_llm_with_retry(
            llm_tool=llm_tool,
            prompt=prompt,
            parser=parser,
            prompt_kwargs={
                "excerpt": excerpt,
                "ontologies_list": ontologies_list,
                "num_ontologies": num_ontologies,
                "format_instructions": parser.get_format_instructions(),
            },
        )

        # Map answer_index to ontology
        # answer_index: 0 -> select None
        # answer_index: 1 to num_ontologies -> select ontology at (answer_index - 1)
        state.status = Status.SUCCESS
        if selector.answer_index == 0:
            # None selected
            logger.debug("LLM selected: None (no suitable ontology)")
            state.current_ontology = NULL_ONTOLOGY
        elif 1 <= selector.answer_index <= num_ontologies:
            # Select ontology at index (answer_index - 1) since list is 0-based
            selected_ontology = ontologies[selector.answer_index - 1]
            logger.debug(
                f"LLM selected ontology at index {selector.answer_index}: "
                f"{selected_ontology.ontology_id} ({selected_ontology.iri})"
            )
            state.current_ontology = selected_ontology
            state.status = Status.SUCCESS
        else:
            # This should not happen due to Pydantic validation, but handle gracefully
            logger.warning(
                f"Invalid answer_index {selector.answer_index} defaulting to NULL_ONTOLOGY"
            )
            state.current_ontology = NULL_ONTOLOGY
    else:
        state.current_ontology = NULL_ONTOLOGY

    # Set the initial version if not already set (tracks original version when ontology was selected)
    if state.current_ontology.initial_version is None:
        state.current_ontology.initial_version = state.current_ontology.version
        logger.debug(
            f"Set initial version for ontology {state.current_ontology.ontology_id}: {state.current_ontology.initial_version}"
        )

    logger.debug(f"Current ontology set to: {state.current_ontology.ontology_id}")

    return state