from __future__ import annotations

from app.core.agent.processes.v2.workflows.doc_update_from_feature_v2.subprocesses.common import SubprocessResult
from app.core.agent.processes.v2.workflows.doc_update_from_feature_v2.subprocesses.common.document_composer import (
    TemplateDocumentComposer,
)
from app.core.agent.processes.v2.workflows.doc_update_from_feature_v2.subprocesses.common.rules_catalog import (
    RulesCatalog,
)
from app.core.agent.processes.v2.workflows.doc_update_from_feature_v2.subprocesses.common.template_registry import (
    TemplateRegistry,
)
from app.core.agent.processes.v2.workflows.doc_update_from_feature_v2.subprocesses.edit_doc.steps.step1_generate_frontmatter.step import (
    EditFrontmatterStep,
)
from app.core.agent.processes.v2.workflows.doc_update_from_feature_v2.subprocesses.edit_doc.steps.step2_generate_sections.step import (
    EditSectionsStep,
)
from app.core.agent.processes.v2.workflows.doc_update_from_feature_v2.subprocesses.edit_doc.steps.step4_compose_document.step import (
    ComposeEditedDocumentStep,
)
from app.core.agent.processes.v2.workflows.doc_update_from_feature_v2.workflow_runtime.models import RequirementTaskContext
from app.core.agent.processes.v2.workflows.doc_update_from_feature_v2.workflow_runtime.models import RuleDocument
from app.core.agent.utils.llm import AgentLlmService


class EditDocSubprocess:
    def __init__(self, llm: AgentLlmService) -> None:
        self._frontmatter = EditFrontmatterStep(llm)
        self._sections = EditSectionsStep(llm)
        self._templates = TemplateRegistry()
        self._compose = ComposeEditedDocumentStep(TemplateDocumentComposer())

    def run(
        self,
        task: RequirementTaskContext,
        *,
        current_content: str,
        rule_documents: list[RuleDocument],
        shared_context: dict[str, object],
    ) -> SubprocessResult:
        catalog = RulesCatalog.from_documents(rule_documents)
        template = self._templates.load(catalog, task.doc_type)
        frontmatter = self._frontmatter.run(
            task,
            current_content=current_content,
            template=template,
            catalog=catalog,
            shared_context=shared_context,
        )
        sections = self._sections.run(
            task,
            current_content=current_content,
            template=template,
            catalog=catalog,
            shared_context=shared_context,
        )
        content = self._compose.run(frontmatter, template, sections, fallback_title=task.heading)
        return SubprocessResult(task=task, content=content)
