from typing import override

from srs2std.generate.standardize_strategies.naive_each_prompt import construct_prompts
from srs2std.generate.standardize_strategies.strategy import StandardizeStrategy
from srs2std.llms import LLM, FinishReason
from srs2std.llms.gemini import Gemini_2_5_Flash, Gemini_2_5_Pro
from srs2std.srs.template import Template


class NaiveEach(StandardizeStrategy):
    def __init__(self, llm: LLM):
        self.llm = llm

    @override
    def run(self, srs_orig: str, template: Template) -> str:
        prompts = construct_prompts(srs_orig, template)
        titles = [
            section_template.title
            for section_template in template.section_templates
            if section_template.description
        ]

        srs_sections = []
        for title, prompt in zip(titles, prompts):
            # request to generate
            response = self.llm.ask(prompt)
            if response.finish_reason != FinishReason.STOP or response.content is None:
                return ""

            # trim response
            section_content = response.content
            section_content = (
                section_content.strip().strip("`").removeprefix("markdown").strip()
            )

            # push front the generated section title if not exists
            first_line = section_content.strip().splitlines()[0].strip().lower()
            if title.number not in first_line or title.name.lower() not in first_line:
                section_content = "\n".join([title.to_str(), section_content])

            srs_sections.append(section_content)

            from time import sleep

            # to avoid rate limit (per minutes)
            model_type = type(self.llm)
            if model_type in [Gemini_2_5_Flash, Gemini_2_5_Pro]:
                sleep(0.4)  # > 60 / 150

        srs_std = "\n\n".join(srs_sections)
        return srs_std
