from __future__ import annotations

import json

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.source_sections import (
    RequirementSourceSections,
)
from app.core.agent.processes.v2.workflows.doc_update_from_feature_v2.subprocesses.common.template_models import (
    SectionDraft,
    TemplateSpec,
)
from app.core.agent.processes.v2.workflows.doc_update_from_feature_v2.workflow_runtime.models import RequirementTaskContext
from app.core.agent.utils.llm import AgentLlmService


class EditSectionsStep:
    def __init__(self, llm: AgentLlmService, source_sections: RequirementSourceSections | None = None) -> None:
        self._llm = llm
        self._source_sections = source_sections or RequirementSourceSections()

    def run(
        self,
        task: RequirementTaskContext,
        *,
        current_content: str,
        template: TemplateSpec,
        catalog: RulesCatalog,
        shared_context: dict[str, object],
    ) -> list[SectionDraft]:
        source_sections = self._source_sections.extract(task.body)
        drafts: list[SectionDraft] = []
        for section in template.sections:
            content = self._generate_section(task, current_content, template, catalog, section.title, section.level, section.rule_path, section.has_children, source_sections, shared_context)
            drafts.append(SectionDraft(title=section.title, level=section.level, content=content))
        return drafts

    def _generate_section(
        self,
        task: RequirementTaskContext,
        current_content: str,
        template: TemplateSpec,
        catalog: RulesCatalog,
        title: str,
        level: int,
        rule_path: str,
        has_children: bool,
        source_sections: dict[str, str],
        shared_context: dict[str, object],
    ) -> str:
        source_fragment = self._source_sections.find_for_title(source_sections, title)
        if has_children and not source_fragment:
            return ""
        payload = {
            "task": {
                "section_key": task.section_key,
                "heading": task.heading,
                "doc_type": task.doc_type,
                "doc_id": task.doc_id,
                "domain": task.domain,
                "sub_domain": task.subdomain,
                "application": task.application,
                "platform": task.platform,
                "path": task.path,
                "requirement_text": task.body,
            },
            "current_content": current_content,
            "template": {
                "doc_type": template.doc_type,
                "source_name": template.source_name,
                "template_text": template.template_text,
                "sections": [{"title": item.title, "level": item.level, "rule_path": item.rule_path} for item in template.sections],
            },
            "target_section": {
                "title": title,
                "level": level,
                "rule_path": rule_path,
                "rule_text": catalog.section_text(rule_path),
            },
            "source_fragment": source_fragment,
            "all_source_sections": source_sections,
            "global_rules": catalog.global_text(),
            "shared_context": shared_context,
        }
        return self._llm.generate(
            "v2_docs_update_v2.edit_section",
            json.dumps(payload, ensure_ascii=False, indent=2),
            log_context=f"workflow.v2.docs_update.from_feature_v2.edit.section.{title}",
        ).strip()
