from pathlib import Path

from srs2std.generate.standardize_context import StandardizeContext
from srs2std.generate.standardize_strategies import NaiveAll, NaiveEach
from srs2std.llms.gemini import Gemini_2_5_Flash, Gemini_2_5_Pro
from srs2std.llms.gpt import GPT_5, GPT_4o
from srs2std.srs.instance import parse_md_as_srs, save_srs_yaml
from srs2std.srs.structure import complement_falling_sections
from srs2std.srs.template import load_template_yaml
from srs2std.util import now_string

LLM_FACTORY = {
    "gemini_2_5_pro": Gemini_2_5_Pro,
    "gemini_2_5_flash": Gemini_2_5_Flash,
    "gpt_5": GPT_5,
    "gpt_4o": GPT_4o,
}

STRATEGY_FACTORY = {
    "naive_all": NaiveAll,
    "naive_each": NaiveEach,
}


def main(
    srs_path: Path,
    std_template_path: Path,
    output_dir: Path,
    prompt_pattern: str,
    llm_name: str,
):
    srs_orig = srs_path.read_text(encoding="utf-8")
    std_template = load_template_yaml(std_template_path)
    std_toc = std_template.to_toc()

    llm_class = LLM_FACTORY.get(llm_name)
    if llm_class is None:
        raise ValueError(f"Unknown LLM name: {llm_name}")
    llm = llm_class()

    strategy_class = STRATEGY_FACTORY.get(prompt_pattern)
    if strategy_class is None:
        raise ValueError(f"Unknown prompt pattern: {prompt_pattern}")
    strategy = strategy_class(llm)

    context = StandardizeContext(srs_orig, std_template)
    srs_std = context.run(strategy=strategy)

    output_dir = output_dir / prompt_pattern / llm_name / now_string()
    output_dir.mkdir(parents=True, exist_ok=True)

    output_path = output_dir / "raw.md"
    output_path.write_text(srs_std, encoding="utf-8")

    srs_std = complement_falling_sections(srs_std, std_toc)
    output_path = output_dir / "postprocessed.md"
    output_path.write_text(srs_std, encoding="utf-8")

    srs_std = parse_md_as_srs(srs_std, std_toc)
    output_path = output_dir / "postprocessed.yaml"
    save_srs_yaml(srs_std, output_path)
