//! Pure template patchback classification policy.

use std::fmt;

use agent_doc_element::element;
use agent_doc_flow::types::{FlowEvent, FlowName, FlowOutcome, FlowStage};

use crate::template::{PatchBlock, parse_patches};

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum PatchbackShape {
    ValidPatch,
    PlainResponse,
    MalformedPatch,
    TranscriptDump,
    MixedOutput,
    /// Raw template form: stdin carries literal `<!-- agent:NAME -->` /
    /// `<!-- /agent:NAME -->` component blocks instead of supported
    /// `<!-- patch:* -->` patch blocks. Committing these as plain text escapes
    /// the markers into the live exchange (`#closeout-repair-churn`), so this
    /// shape must fail closed before commit.
    EscapedComponentMarkers,
}

impl PatchbackShape {
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::ValidPatch => "valid_patch",
            Self::PlainResponse => "plain_response",
            Self::MalformedPatch => "malformed_patch",
            Self::TranscriptDump => "transcript_dump",
            Self::MixedOutput => "mixed_output",
            Self::EscapedComponentMarkers => "escaped_component_markers",
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PatchbackShapeFacts {
    pub marker_count: usize,
    pub patch_count: usize,
    pub exchange_patch_count: usize,
    pub unmatched_len: usize,
    pub has_transcript_markers: bool,
    /// Count of fully-formed `<!-- agent:NAME -->` ... `<!-- /agent:NAME -->`
    /// component blocks parsed from the response stdin (outside code fences).
    /// Non-zero with no patch markers means the caller piped the raw template
    /// form instead of `<!-- patch:* -->` blocks.
    pub raw_component_block_count: usize,
}

pub fn classify_patchback_shape(facts: PatchbackShapeFacts) -> PatchbackShape {
    if facts.marker_count > 0 && facts.patch_count == 0 {
        return PatchbackShape::MalformedPatch;
    }
    if facts.patch_count > 0 && facts.unmatched_len > 0 {
        return PatchbackShape::MixedOutput;
    }
    if facts.patch_count > 0 {
        return PatchbackShape::ValidPatch;
    }
    // No supported patch blocks. A response that instead carries literal
    // component blocks is the raw template form and must not be committed as
    // plain exchange text (`#closeout-repair-churn`).
    if facts.raw_component_block_count > 0 {
        return PatchbackShape::EscapedComponentMarkers;
    }
    if facts.has_transcript_markers {
        return PatchbackShape::TranscriptDump;
    }
    PatchbackShape::PlainResponse
}

/// Count fully-formed `<!-- agent:NAME -->` component blocks in a response.
///
/// Reuses the validated component parser so boundary markers
/// (`<!-- agent:boundary:HASH -->`, which have no close marker) and markers
/// inside code fences are excluded. A parse error (mismatched/invalid markers)
/// is treated as "no clean component blocks" because the malformed/plain paths
/// handle that shape separately.
pub fn raw_component_block_count(response: &str) -> usize {
    element::parse(response)
        .map(|components| components.len())
        .unwrap_or(0)
}

pub fn patchback_marker_count_outside_code(response: &str) -> usize {
    let code_ranges = element::find_code_ranges(response);
    let markers = [
        "<!-- patch:",
        "<!-- /patch:",
        "<!-- replace:",
        "<!-- /replace:",
    ];
    let mut count = 0usize;
    for marker in markers {
        let mut search_from = 0usize;
        while let Some(rel) = response[search_from..].find(marker) {
            let pos = search_from + rel;
            if !code_ranges
                .iter()
                .any(|&(start, end)| pos >= start && pos < end)
            {
                count += 1;
            }
            search_from = pos + marker.len();
        }
    }
    count
}

pub fn patchback_shape_facts(
    response: &str,
    patches: &[PatchBlock],
    unmatched: &str,
) -> PatchbackShapeFacts {
    PatchbackShapeFacts {
        marker_count: patchback_marker_count_outside_code(response),
        patch_count: patches.len(),
        exchange_patch_count: patches
            .iter()
            .filter(|patch| patch.name == "exchange")
            .count(),
        unmatched_len: unmatched.trim().len(),
        has_transcript_markers: response.contains("## User") || response.contains("## Assistant"),
        raw_component_block_count: raw_component_block_count(response),
    }
}

#[derive(Debug, Clone)]
pub struct TemplatePatchbackPlan {
    pub patches: Vec<PatchBlock>,
    pub unmatched: String,
    pub shape: PatchbackShape,
    pub marker_count: usize,
    pub exchange_patch_count: usize,
    pub raw_component_block_count: usize,
}

impl TemplatePatchbackPlan {
    pub fn has_malformed_orphan_markers(&self) -> bool {
        self.shape == PatchbackShape::MalformedPatch && !self.unmatched.trim().is_empty()
    }

    pub fn has_escaped_component_markers(&self) -> bool {
        self.shape == PatchbackShape::EscapedComponentMarkers
    }
}

pub fn parse_template_patchback_plan(response: &str) -> anyhow::Result<TemplatePatchbackPlan> {
    let (patches, unmatched) = parse_patches(response)?;
    let facts = patchback_shape_facts(response, &patches, &unmatched);
    let shape = classify_patchback_shape(facts.clone());
    Ok(TemplatePatchbackPlan {
        patches,
        unmatched,
        shape,
        marker_count: facts.marker_count,
        exchange_patch_count: facts.exchange_patch_count,
        raw_component_block_count: facts.raw_component_block_count,
    })
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OrchestratePatchbackRejectReason {
    MissingExchangePatch,
    RawUnmatchedContent,
    FullDocumentComponentMarkers,
    TranscriptPromptLines,
    TranscriptHeadings,
    MultipleAssistantResponses,
}

impl OrchestratePatchbackRejectReason {
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::MissingExchangePatch => "missing_exchange_patch",
            Self::RawUnmatchedContent => "raw_unmatched_content",
            Self::FullDocumentComponentMarkers => "full_document_component_markers",
            Self::TranscriptPromptLines => "transcript_prompt_lines",
            Self::TranscriptHeadings => "transcript_headings",
            Self::MultipleAssistantResponses => "multiple_assistant_responses",
        }
    }

    pub const fn message(self) -> &'static str {
        match self {
            Self::MissingExchangePatch => {
                "orchestrate template-mode responses must include a <!-- patch:exchange --> block"
            }
            Self::RawUnmatchedContent => {
                "orchestrate template-mode responses must not include raw unmatched content outside patch blocks"
            }
            Self::FullDocumentComponentMarkers => {
                "orchestrate template-mode plain responses must not include full document component markers"
            }
            Self::TranscriptPromptLines => {
                "orchestrate template-mode plain responses must not include transcript prompt lines"
            }
            Self::TranscriptHeadings => {
                "orchestrate template-mode plain responses must not include transcript headings"
            }
            Self::MultipleAssistantResponses => {
                "orchestrate template-mode plain responses must contain only one assistant response"
            }
        }
    }
}

impl fmt::Display for OrchestratePatchbackRejectReason {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.message())
    }
}

impl std::error::Error for OrchestratePatchbackRejectReason {}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OrchestratePatchbackDecision {
    AcceptExplicitPatch,
    AcceptPlainResponse,
    Reject(OrchestratePatchbackRejectReason),
}

impl OrchestratePatchbackDecision {
    pub fn is_accepted(self) -> bool {
        !matches!(self, Self::Reject(_))
    }
}

pub fn classify_orchestrate_patchback(
    patches: &[PatchBlock],
    unmatched: &str,
) -> OrchestratePatchbackDecision {
    if patches.is_empty() {
        return classify_orchestrate_plain_response(unmatched);
    }
    if !patches.iter().any(|patch| patch.name == "exchange") {
        return OrchestratePatchbackDecision::Reject(
            OrchestratePatchbackRejectReason::MissingExchangePatch,
        );
    }
    if !unmatched.trim().is_empty() {
        return OrchestratePatchbackDecision::Reject(
            OrchestratePatchbackRejectReason::RawUnmatchedContent,
        );
    }
    OrchestratePatchbackDecision::AcceptExplicitPatch
}

pub fn classify_orchestrate_plain_response(unmatched: &str) -> OrchestratePatchbackDecision {
    let trimmed = unmatched.trim();
    if trimmed.is_empty() {
        return OrchestratePatchbackDecision::AcceptPlainResponse;
    }

    if trimmed.contains("<!-- agent:")
        || trimmed.contains("<!-- /agent:")
        || trimmed.contains("&lt;!-- agent:")
        || trimmed.contains("&lt;!-- /agent:")
    {
        return OrchestratePatchbackDecision::Reject(
            OrchestratePatchbackRejectReason::FullDocumentComponentMarkers,
        );
    }

    if trimmed
        .lines()
        .any(|line| line.trim_start().starts_with('❯'))
    {
        return OrchestratePatchbackDecision::Reject(
            OrchestratePatchbackRejectReason::TranscriptPromptLines,
        );
    }

    if trimmed.lines().any(|line| {
        let line = line.trim();
        line == "## User"
            || line.starts_with("## User ")
            || line == "## Assistant"
            || line.starts_with("## Assistant ")
    }) {
        return OrchestratePatchbackDecision::Reject(
            OrchestratePatchbackRejectReason::TranscriptHeadings,
        );
    }

    let response_headings = trimmed
        .lines()
        .filter(|line| line.trim_start().starts_with("### Re:"))
        .count();
    if response_headings > 1 {
        return OrchestratePatchbackDecision::Reject(
            OrchestratePatchbackRejectReason::MultipleAssistantResponses,
        );
    }

    OrchestratePatchbackDecision::AcceptPlainResponse
}

pub fn enforce_orchestrate_patchback_contract(
    origin: Option<&str>,
    patches: &[PatchBlock],
    unmatched: &str,
) -> Result<(), OrchestratePatchbackRejectReason> {
    if origin != Some("orchestrate") {
        return Ok(());
    }
    match classify_orchestrate_patchback(patches, unmatched) {
        OrchestratePatchbackDecision::AcceptExplicitPatch
        | OrchestratePatchbackDecision::AcceptPlainResponse => Ok(()),
        OrchestratePatchbackDecision::Reject(reason) => Err(reason),
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ChildPatchbackNormalizationDecision {
    WrappedPlainResponse,
    KeptExplicitPatch,
    KeptRejectedPlainResponse,
    KeptUnparseable,
}

impl ChildPatchbackNormalizationDecision {
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::WrappedPlainResponse => "wrapped_plain_response",
            Self::KeptExplicitPatch => "kept_explicit_patch",
            Self::KeptRejectedPlainResponse => "kept_rejected_plain_response",
            Self::KeptUnparseable => "kept_unparseable",
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ChildPatchbackNormalization {
    pub response: String,
    pub decision: ChildPatchbackNormalizationDecision,
}

pub fn normalize_child_template_response(response: String) -> ChildPatchbackNormalization {
    let Ok((patches, unmatched)) = parse_patches(&response) else {
        return ChildPatchbackNormalization {
            response,
            decision: ChildPatchbackNormalizationDecision::KeptUnparseable,
        };
    };
    if patches.is_empty() && !unmatched.trim().is_empty() {
        if classify_orchestrate_plain_response(unmatched.trim()).is_accepted() {
            return ChildPatchbackNormalization {
                response: format!(
                    "<!-- patch:exchange -->\n{}\n<!-- /patch:exchange -->\n",
                    unmatched.trim_end()
                ),
                decision: ChildPatchbackNormalizationDecision::WrappedPlainResponse,
            };
        }
        return ChildPatchbackNormalization {
            response,
            decision: ChildPatchbackNormalizationDecision::KeptRejectedPlainResponse,
        };
    }

    let decision = if matches!(
        classify_orchestrate_patchback(&patches, &unmatched),
        OrchestratePatchbackDecision::AcceptExplicitPatch
    ) {
        ChildPatchbackNormalizationDecision::KeptExplicitPatch
    } else {
        ChildPatchbackNormalizationDecision::KeptRejectedPlainResponse
    };
    ChildPatchbackNormalization { response, decision }
}

pub fn child_patchback_normalization_event(
    normalization: &ChildPatchbackNormalization,
) -> FlowEvent {
    let outcome = match normalization.decision {
        ChildPatchbackNormalizationDecision::WrappedPlainResponse
        | ChildPatchbackNormalizationDecision::KeptExplicitPatch => FlowOutcome::Completed,
        ChildPatchbackNormalizationDecision::KeptRejectedPlainResponse
        | ChildPatchbackNormalizationDecision::KeptUnparseable => FlowOutcome::FailedClosed,
    };
    FlowEvent::new(
        FlowName::OrchestrationBatch,
        FlowStage::ChildCloseout,
        outcome,
    )
    .with_reason(format!(
        "child_patchback:{}",
        normalization.decision.as_str()
    ))
}

pub fn child_template_finalize_text(response: String) -> String {
    normalize_child_template_response(response).response
}

pub fn should_stream_exchange_patch(response: &str) -> bool {
    response.contains("<!-- patch:exchange")
}

pub fn finalize_suffix_from_streamed_prefix(streamed: &str, full: &str) -> Option<String> {
    if let Some(delta) = finalize_suffix_from_open_patch_prefix(streamed, full) {
        return Some(delta);
    }

    if let (Ok((full_patches, full_unmatched)), Ok((streamed_patches, streamed_unmatched))) =
        (parse_patches(full), parse_patches(streamed))
        && full_unmatched.trim().is_empty()
        && streamed_unmatched.trim().is_empty()
        && full_patches.len() == streamed_patches.len()
    {
        let mut delta = String::new();
        for (full_patch, streamed_patch) in full_patches.iter().zip(streamed_patches.iter()) {
            if full_patch.name != streamed_patch.name
                || !full_patch.content.starts_with(&streamed_patch.content)
            {
                return None;
            }
            let suffix = &full_patch.content[streamed_patch.content.len()..];
            if suffix.is_empty() {
                continue;
            }
            delta.push_str(&format!(
                "<!-- patch:{} -->\n{}<!-- /patch:{} -->\n",
                full_patch.name, suffix, full_patch.name
            ));
        }

        if !delta.trim().is_empty() {
            return Some(delta);
        }
    }
    None
}

fn finalize_suffix_from_open_patch_prefix(streamed: &str, full: &str) -> Option<String> {
    if !full.starts_with(streamed) {
        return None;
    }

    let open_start = streamed.find("<!-- patch:")?;
    let open_end = streamed[open_start..].find("-->")? + open_start + 3;
    let open_marker = &streamed[open_start..open_end];
    let patch_name = open_marker
        .strip_prefix("<!-- patch:")?
        .strip_suffix(" -->")?
        .trim();
    if patch_name.is_empty() {
        return None;
    }

    let mut content_start = open_end;
    if streamed.as_bytes().get(content_start) == Some(&b'\n') {
        content_start += 1;
    }
    let close_marker = format!("<!-- /patch:{} -->", patch_name);
    let close_pos = full[content_start..].find(&close_marker)? + content_start;
    let full_content = &full[content_start..close_pos];
    let streamed_content = &streamed[content_start..];
    if !full_content.starts_with(streamed_content) {
        return None;
    }
    let suffix = &full_content[streamed_content.len()..];
    if suffix.is_empty() {
        return None;
    }

    Some(format!(
        "<!-- patch:{} -->\n{}<!-- /patch:{} -->\n",
        patch_name, suffix, patch_name
    ))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn patch_markers_without_closed_blocks_are_malformed() {
        let shape = classify_patchback_shape(PatchbackShapeFacts {
            marker_count: 2,
            patch_count: 0,
            exchange_patch_count: 0,
            unmatched_len: 42,
            has_transcript_markers: false,
            raw_component_block_count: 0,
        });

        assert_eq!(shape, PatchbackShape::MalformedPatch);
    }

    #[test]
    fn patch_plus_unmatched_text_is_mixed_output() {
        let shape = classify_patchback_shape(PatchbackShapeFacts {
            marker_count: 1,
            patch_count: 1,
            exchange_patch_count: 1,
            unmatched_len: 7,
            has_transcript_markers: false,
            raw_component_block_count: 0,
        });

        assert_eq!(shape, PatchbackShape::MixedOutput);
    }

    #[test]
    fn raw_component_blocks_without_patches_are_escaped_component_markers() {
        let shape = classify_patchback_shape(PatchbackShapeFacts {
            marker_count: 0,
            patch_count: 0,
            exchange_patch_count: 0,
            unmatched_len: 64,
            has_transcript_markers: false,
            raw_component_block_count: 1,
        });

        assert_eq!(shape, PatchbackShape::EscapedComponentMarkers);
    }

    #[test]
    fn valid_patches_take_precedence_over_incidental_component_blocks() {
        let shape = classify_patchback_shape(PatchbackShapeFacts {
            marker_count: 1,
            patch_count: 1,
            exchange_patch_count: 1,
            unmatched_len: 0,
            has_transcript_markers: false,
            raw_component_block_count: 1,
        });

        assert_eq!(shape, PatchbackShape::ValidPatch);
    }

    #[test]
    fn raw_component_block_count_ignores_plain_response_and_boundary_markers() {
        assert_eq!(
            raw_component_block_count("### Re: topic - gpt-5\n\nImplemented and verified.\n"),
            0
        );
        assert_eq!(
            raw_component_block_count("<!-- agent:boundary:abc123 -->\nsome text\n"),
            0
        );
        assert_eq!(
            raw_component_block_count(
                "<!-- agent:exchange -->\n### Re: topic\nbody\n<!-- /agent:exchange -->\n"
            ),
            1
        );
    }

    #[test]
    fn parse_template_patchback_plan_accepts_plain_response_without_component_markers() {
        let plain = "### Re: closeout - gpt-5\n\nImplemented and verified.\n";
        let plan = parse_template_patchback_plan(plain).unwrap();
        assert_eq!(plan.shape, PatchbackShape::PlainResponse);
    }

    #[test]
    fn orchestrate_plain_response_rejects_transcript_prompt_lines() {
        let decision = classify_orchestrate_plain_response("### Re: done\n\nDone.\n❯ do #next\n");

        assert_eq!(
            decision,
            OrchestratePatchbackDecision::Reject(
                OrchestratePatchbackRejectReason::TranscriptPromptLines
            )
        );
    }

    #[test]
    fn orchestrate_patchback_requires_exchange_patch() {
        let patches = vec![PatchBlock::new("status", "updated")];
        let decision = classify_orchestrate_patchback(&patches, "");

        assert_eq!(
            decision,
            OrchestratePatchbackDecision::Reject(
                OrchestratePatchbackRejectReason::MissingExchangePatch
            )
        );
    }

    #[test]
    fn child_plain_response_is_wrapped_as_exchange_patch() {
        let normalized =
            normalize_child_template_response("### Re: child — gpt-5\n\nImplemented.".to_string());

        assert_eq!(
            normalized.decision,
            ChildPatchbackNormalizationDecision::WrappedPlainResponse
        );
        assert!(normalized.response.starts_with("<!-- patch:exchange -->"));
        assert!(normalized.response.contains("### Re: child"));
    }

    #[test]
    fn child_patchback_normalization_event_marks_wrapped_plain_response_completed() {
        let normalized =
            normalize_child_template_response("### Re: child — gpt-5\n\nImplemented.".to_string());
        let event = child_patchback_normalization_event(&normalized);

        assert_eq!(event.flow, FlowName::OrchestrationBatch);
        assert_eq!(event.stage, FlowStage::ChildCloseout);
        assert_eq!(event.outcome, FlowOutcome::Completed);
        assert_eq!(
            event.reason.as_deref(),
            Some("child_patchback:wrapped_plain_response")
        );
    }

    #[test]
    fn transcript_shaped_child_response_is_not_normalized() {
        let transcript = "❯ do #next\n### Re: malformed — gpt-5\nBody";
        let normalized = normalize_child_template_response(transcript.to_string());

        assert_eq!(
            normalized.decision,
            ChildPatchbackNormalizationDecision::KeptRejectedPlainResponse
        );
        assert_eq!(normalized.response, transcript);
    }

    #[test]
    fn finalize_suffix_uses_only_unseen_stream_tail() {
        let streamed = "<!-- patch:exchange -->\n### Re: streamed - gpt-5\n";
        let full = "<!-- patch:exchange -->\n### Re: streamed - gpt-5\n\nImplemented and verified.\n<!-- /patch:exchange -->\n";
        let delta = finalize_suffix_from_streamed_prefix(streamed, full).unwrap();
        assert!(!delta.contains("### Re: streamed - gpt-5"));
        assert!(delta.contains("Implemented and verified."));
    }

    #[test]
    fn child_template_finalize_wraps_plain_response_as_exchange_patch() {
        let finalize = child_template_finalize_text(
            "### Re: plain orch - gpt-5\n\nImplemented and verified.".to_string(),
        );

        assert!(finalize.starts_with("<!-- patch:exchange -->"));
        assert!(finalize.contains("### Re: plain orch"));
        assert!(finalize.ends_with("<!-- /patch:exchange -->\n"));
    }

    #[test]
    fn child_template_finalize_does_not_wrap_transcript_response() {
        let transcript = "❯ do #next\n### Re: malformed - gpt-5\nBody";
        let finalize = child_template_finalize_text(transcript.to_string());

        assert_eq!(finalize, transcript);
    }

    #[test]
    fn streamed_flush_waits_for_exchange_patch_marker() {
        assert!(!should_stream_exchange_patch(
            "### Re: malformed streaming closeout - gpt-5\nBody"
        ));
        assert!(should_stream_exchange_patch(
            "<!-- patch:exchange -->\n### Re: streamed - gpt-5\n"
        ));
    }
}
