# -*- coding: utf-8 -*-
"""
Darwin-60B-DUO Sequential Refine — two-model collaboration.

drafter_backend produces the initial draft, then refiner_backend polishes it.
The polish prompt is built dynamically based on the language combination so
that:
  - Darwin (English reasoning) → AWAXIS (Korean polish) for Korean output
    requiring rigorous English/STEM reasoning
  - AWAXIS (Korean cultural context) → Darwin (English polish) for English
    output requiring Korean cultural / linguistic context
"""
import re
from typing import Any, Dict, List


def _last_user_text(messages: List[Dict[str, str]]) -> str:
    for m in reversed(messages):
        if m.get("role") == "user":
            return m.get("content", "")
    return ""


def _korean_ratio(text: str) -> float:
    if not text:
        return 0.0
    return len(re.findall(r"[가-힣]", text)) / len(text)


async def sequential_refine(
    drafter,
    refiner,
    messages: List[Dict[str, str]],
    temperature: float = 0.5,
    max_tokens: int = 4096,
) -> str:
    """
    Step 1: drafter produces the initial answer using the user's messages.
    Step 2: refiner is given the original messages + the drafter's response +
            a polish instruction, then produces the final output.

    The polish instruction is language-adaptive:
      - If user asked in Korean (kr_ratio > 0.3) → polish to natural Korean
      - If user asked in English → polish to clearer English
      - Otherwise → general clarity polish
    """
    user_text = _last_user_text(messages)
    kr = _korean_ratio(user_text)

    # ---- Step 1: drafter ----
    draft_outputs = await drafter.chat(
        messages,
        temperature=temperature,
        max_tokens=max_tokens,
    )
    draft = draft_outputs[0]

    # ---- Step 2: refiner polish ----
    if kr > 0.3:
        polish_instruction = (
            "위 초안을 사용자의 원래 질문 의도에 맞게 한국어로 자연스럽고 "
            "정확하게 다듬어 최종 답변을 작성하세요. 사실관계는 보존하되, "
            "어색한 표현·번역체·중복은 제거하고, 한국어 독자에게 매끄러운 "
            "흐름이 되도록 재작성하세요. 새로운 정보 추가 금지 — 표현만 정련하세요."
        )
    elif kr < 0.05 and len(user_text) > 0:
        polish_instruction = (
            "Polish the draft above into a clearer, more concise, and "
            "natural-sounding English response that fully addresses the "
            "user's original question. Preserve all factual content; remove "
            "redundancy, awkward phrasing, and translation artifacts. Do "
            "not add new information — refine wording only."
        )
    else:
        polish_instruction = (
            "Refine the draft above for clarity, naturalness, and "
            "consistency. Preserve all facts; remove redundancy. Do not "
            "introduce new information."
        )

    refine_messages = list(messages) + [
        {"role": "assistant", "content": draft},
        {"role": "user", "content": polish_instruction},
    ]
    refined_outputs = await refiner.chat(
        refine_messages,
        temperature=max(0.0, temperature - 0.2),  # cooler for polish
        max_tokens=max_tokens,
    )
    return refined_outputs[0]
