import os

import instructor
from dotenv import dotenv_values
from pydantic import BaseModel, Field
from vertexai.generative_models import GenerativeModel

config = dotenv_values(".env")
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = config["GOOGLE_APPLICATION_CREDENTIALS"]
client = instructor.from_vertexai(
    client=GenerativeModel("gemini-1.5-flash-002"),
    mode=instructor.Mode.VERTEXAI_TOOLS,
)

last_line_number = 2
annotated_content = "Line 0: a\nLine 1: b\nLine 2: c"
system_prompt = f"""
You are an expert web content analyzer. Your task is to analyze the provided webpage text and generate 3 topic headings that cover the full content of the webpage. 
For each topic heading, you must also provide the starting and ending line numbers of the corresponding section. Each section is the continuation of the previous section.
The line numbers must be integers, start from `0`, and end at the last line `{last_line_number}`. There should be no overlap between the line number ranges.
""".strip()


class TopicHeading(BaseModel):
    reasoning: str = Field(
        description="The reasoning process of why the topic heading and line numbers are chosen."
    )
    heading: str = Field(description="The topic heading for a section of the webpage.")
    start_line: int = Field(description="The starting line number of the section.")
    end_line: int = Field(description="The ending line number of the section.")


class WebpageTopics(BaseModel):
    topics: list[TopicHeading] = Field(
        description="A list of topic headings with their corresponding start and end line numbers."
    )


resp, completion = client.chat.completions.create_with_completion(
    messages=[
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": annotated_content},
    ],
    response_model=WebpageTopics
)

print(resp)