"""
Agent module that performs structured extraction using LangChain and Pydantic.

The module exposes a single public function `extract_structured_output` which:
1. Detects whether the input text describes a person or a meeting.
2. Builds an appropriate prompt with format instructions from the chosen parser.
3. Sends the request to the LLM via LangChain and returns the parsed Pydantic model instance.
"""

from __future__ import annotations

import re
from typing import Union, List

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

from models import PersonInfo, MeetingNotes

# LLM configuration – BroJS endpoint
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
    api_key=None,  # will be taken from env var JOURNAL_MCP_PAT
    temperature=0.0,
)

# Prompt template – the same for both schemas; format instructions are injected.
PROMPT = PromptTemplate(
    input_variables=["text", "format_instructions"],
    template="""
You are a data extraction assistant.
Extract structured information from the following text and return it as JSON that matches the provided schema.

Text:
{text}

Schema format instructions:
{format_instructions}

Respond only with valid JSON. Do not add any extra text.
""",
)

def _detect_schema(text: str) -> type:
    """Heuristically determine whether the input describes a person or a meeting.
    Returns the corresponding Pydantic model class.
    """
    # Simple keyword checks – can be extended
    if re.search(r"\bmeeting\b|\bdate\b|\bparticipants?\b", text, re.I):
        return MeetingNotes
    return PersonInfo


def extract_structured_output(text: str) -> Union[PersonInfo, MeetingNotes]:
    """Return a Pydantic model instance extracted from *text*.

    Parameters
    ----------
    text:
        Free‑form input string.

    Returns
    -------
    PersonInfo | MeetingNotes
        Parsed data as a validated Pydantic object.
    """
    schema = _detect_schema(text)
    parser = PydanticOutputParser(pydantic_object=schema)
    prompt = PROMPT.format(text=text, format_instructions=parser.get_format_instructions())
    response = llm.invoke([prompt])
    # The LLM returns a string; parse it
    return parser.parse(response.content)

__all__ = ["extract_structured_output"]