from pydantic import BaseModel, Field
from typing import List, Optional, Dict
from pydantic_ai import Agent

from llama_index.core.program import LLMTextCompletionProgram
from llama_index.core.extractors import PydanticProgramExtractor
from llama_index.core.output_parsers import PydanticOutputParser

from ketju.models import ChapterMetadata


agent: Agent[None, ChapterMetadata] = Agent[None, ChapterMetadata](
    model="openai:gpt-4o-mini",
    system_prompt="""
        You are analyzing a section of a knowledge document.

        Task:
        - Produce a clear, accurate summary of the content
        - Identify key concepts, terms, and topics covered
        - Generate natural questions that readers might ask
        - Note any cross-references to other materials

        Rules:
        - Do NOT reference "the document" or "this section" in the summary
        - Do NOT include questions in the summary
        - Questions must be concise and specific to the content
        - Focus on factual content and main ideas
        - Extract key concepts that define the topic

        """,
    output_type=ChapterMetadata,
)

def create_chapter_metadata(context: str) -> ChapterMetadata:
    """Create metadata for a document section using the agent."""
    return agent.run_sync(f"Generate chapter metadata for the following text: {context}")

chapter_metadata_program = LLMTextCompletionProgram.from_defaults(
    output_parser=PydanticOutputParser(output_cls=ChapterMetadata),
    prompt_template_str="""
        You are analyzing a section of a knowledge document.

        Task:
        - Produce a clear, accurate summary of the content
        - Identify key concepts, terms, and topics covered
        - Generate natural questions that readers might ask about this content
        - Note any cross-references to other materials or sources

        Rules:
        - Do NOT reference "the document" or "this section" in the summary
        - Do NOT include questions in the summary
        - Questions must be concise and specific to the content
        - Focus on factual content and main ideas
        - Extract key concepts that define the topic

        Text:
        {context}
        """,
)