"""`llm` 스텝 -- 체인 텍스트(`ctx.text`)를 LLM 한 번(또는 self-consistency N번)으로 변환.

`multi_step.py` 의 llm 분기에서 로직 이동 (2026-08-10, 동작 동일).
bulk 분할 · JSON output 추출 · fallback · repeat/intersection(self-consistency) ·
corrections 캡처를 담당한다.
"""

import json
import logging
import re
from datetime import datetime
from typing import Any

from solar_eval.models.sample import EvalSample
from solar_eval.pipelines import judge
from solar_eval.pipelines.pipeline import BaseStep, StepContext
from solar_eval.pipelines.registry import register_step
from solar_eval.pipelines.steps.chunking import (
    ChunkPlan,
    join_chunks,
    normalize_separator_literals,
    plan_chunks,
)

logger = logging.getLogger(__name__)

# 출력 길이 가드 -- 교열은 길이를 거의 바꾸지 않는다. 입력 대비 ±25% 또는 40자(짧은 입력용) 를 넘는 변동은
# 모델이 문장을 지어냈거나(실측 49자→321자) bulk 를 통째로 버린 것(490자→7자)이라 폴백/입력 유지로 처리한다.
# 종전 규칙("5배 초과")은 길이 급감과 5배 미만의 환각을 그대로 통과시켰다 (insight 260910).
LENGTH_TOL_REL = 0.25
LENGTH_TOL_ABS = 40


def plausible_length(source: str, output: str) -> bool:
    if not source:
        return True
    return abs(len(output) - len(source)) <= max(LENGTH_TOL_ABS, LENGTH_TOL_REL * len(source))


@register_step("llm")
class LLMStep(BaseStep):
    """체인 텍스트 위에서 도는 LLM 스텝 (prod/v24 파이프라인의 기본 부품)."""

    async def run(self, sample: EvalSample, ctx: StepContext) -> None:
        step = self.config
        step_name = self.name
        step_model = step.get("model", ctx.model)
        step_re = step.get("reasoning_effort", ctx.reasoning_effort)
        step_max_tokens = step.get("max_tokens", ctx.max_tokens)
        step_temperature = step.get("temperature", ctx.temperature)
        step_response_format = step.get("response_format")
        prompt_key = step.get("prompt_key", step_name)
        step_prompt = self.env.prompts.get(prompt_key, step.get("prompt", ""))
        # v24 self-consistency (basic_correction): repeat>1 runs the bulk N
        # times at different temperatures and keeps only corrections seen in
        # >= intersection_threshold runs.
        step_repeat = int(step.get("repeat", 1) or 1)

        include_original = step.get("include_original", False)
        user_template = step.get("user_template")

        text = ctx.text
        original_text = ctx.original_text
        ts = datetime.now().strftime("%H:%M:%S")

        # Split into bulks if needed
        # For include_original/user_template steps, skip splitting
        if include_original or user_template:
            plan = ChunkPlan(chunks=[text], seps=[""], kinds=[], lead="")
        else:
            plan = plan_chunks(
                text, self.env.text_split_size, strategy=self.env.text_split_strategy
            )
        bulks = plan.chunks
        processed_bulks = []
        logger.debug(f"  [{ts}] {len(bulks)} bulk(s), text_len={len(text)}")
        if plan.is_split:
            # 절단면 진단용 -- 어디를 어떤 경계로 잘랐는지 (insight 260910 의 GLUED/LOST 분석이 읽는다)
            chunk_plans = sample.artifacts.setdefault("chunk_plans", {})
            chunk_plans[step_name] = {
                "lens": [len(b) for b in bulks],
                "kinds": list(plan.kinds),
            }

        for bulk_i, bulk in enumerate(bulks):
            if ctx.on_step and len(bulks) > 1:
                ctx.on_step("chunk", bulk_i + 1, len(bulks), step_name, ctx.doc_chunk)
            if step_repeat > 1:
                (
                    final_text,
                    run_usage,
                    runs,
                    survivors,
                ) = await self._run_self_consistency_bulk(
                    bulk,
                    step,
                    step_model=step_model,
                    step_re=step_re,
                    step_max_tokens=step_max_tokens,
                    step_response_format=step_response_format,
                    step_prompt=step_prompt,
                    user_template=user_template,
                    include_original=include_original,
                    original_text=original_text,
                    step_outputs=ctx.step_outputs,
                    ctx=ctx,
                )
                processed_bulks.append(final_text)
                ctx.add_usage(run_usage)
                # 교집합 이전 run 별 제안도 남긴다 -- 진단 문서가 요구한 가시성
                # (§5-F 지시 4). 교집합 결과(survivors)는 downstream
                # tool_calling_judge/apply_judge_decisions 가 "corrections" 로 읽는다.
                sc_runs = sample.artifacts.setdefault("self_consistency_runs", {})
                sc_runs.setdefault(step_name, []).extend(runs)
                step_corrs = sample.artifacts.setdefault("corrections", {})
                step_corrs.setdefault(step_name, []).extend(survivors)
                continue

            if user_template:
                user_content = self._format_user_template(
                    user_template,
                    original_text,
                    bulk,
                    ctx.step_outputs,
                )
            elif include_original:
                user_content = f"[원문]\n{original_text}\n\n[교정본]\n{bulk}"
            else:
                user_content = bulk

            ts2 = datetime.now().strftime("%H:%M:%S")
            logger.debug(
                f"  [{ts2}] bulk {bulk_i + 1}/{len(bulks)} calling API (len={len(bulk)})..."
            )
            msgs: list[dict[str, str]] = []
            if step_prompt:
                msgs.append({"role": "system", "content": step_prompt})
            msgs.append({"role": "user", "content": user_content})
            result = await ctx.provider.chat_completion(
                messages=msgs,
                model=step_model,
                temperature=step_temperature,
                max_tokens=step_max_tokens,
                reasoning_effort=step_re,
                response_format=step_response_format,
            )
            ts3 = datetime.now().strftime("%H:%M:%S")
            raw_len = len(result.get("content", ""))
            output = normalize_separator_literals(
                self._extract_output(result["content"]).rstrip(), source=bulk
            )

            # 길이 가드: JSON 파싱 실패(설명이 통째로 output 이 됨) · 환각 · bulk 삭제를 한 규칙으로 잡는다
            if not plausible_length(bulk, output):
                raw_content = result.get("content", "")
                logger.warning(
                    f"  [{ts3}] [{step_name}] bulk {bulk_i + 1} output length implausible "
                    f"(in={len(bulk)}, out={len(output)}), attempting fallback"
                )
                guard_log = sample.artifacts.setdefault("length_guard", [])
                guard_log.append(
                    {"step": step_name, "bulk": bulk_i + 1, "in": len(bulk), "out": len(output)}
                )

                # Try fallback LLM call if configured
                fallback_config = step.get("fallback")
                fallback_prompt = (
                    self.env.prompts.get(fallback_config["prompt_key"], "")
                    if fallback_config
                    else ""
                )
                if fallback_prompt:
                    fallback_user = (
                        f"<원문>\n{original_text}\n</원문>\n\n"
                        f"<교정_결과_JSON>\n{raw_content}\n</교정_결과_JSON>"
                    )
                    fallback_msgs: list[dict[str, str]] = [
                        {"role": "system", "content": fallback_prompt},
                        {"role": "user", "content": fallback_user},
                    ]
                    fallback_rf = fallback_config.get("response_format")
                    fallback_result = await ctx.provider.chat_completion(
                        messages=fallback_msgs,
                        model=step_model,
                        temperature=step_temperature,
                        max_tokens=step_max_tokens,
                        reasoning_effort=step_re,
                        response_format=fallback_rf,
                    )
                    fallback_output = normalize_separator_literals(
                        self._extract_output(fallback_result["content"]).rstrip(), source=bulk
                    )
                    ts4 = datetime.now().strftime("%H:%M:%S")

                    if plausible_length(bulk, fallback_output):
                        output = fallback_output
                        logger.info(
                            f"  [{ts4}] [{step_name}] fallback succeeded, len={len(output)}"
                        )
                        ctx.add_usage(fallback_result.get("usage", {}))
                    else:
                        output = bulk
                        logger.warning(
                            f"  [{ts4}] [{step_name}] fallback length also implausible, using input text"
                        )
                else:
                    output = bulk
                    logger.warning(
                        f"  [{ts3}] [{step_name}] no fallback configured, using input text"
                    )

            logger.debug(
                f"  [{ts3}] bulk {bulk_i + 1} done, raw={raw_len}, extracted={len(output)}"
            )
            processed_bulks.append(output)
            ctx.add_usage(result.get("usage", {}))

            # response_format 이 "corrections" 배열을 담고 있으면(예:
            # basic_correction) repeat 없이도 뽑아 둔다 -- downstream
            # tool_calling_judge/apply_judge_decisions 가 읽을 수 있게.
            # output 계산과 무관한 순수 추가이므로 기존 텍스트 동작은 바뀌지 않는다.
            step_corrections = judge.extract_corrections(result.get("content", ""))
            if step_corrections:
                step_corrs = sample.artifacts.setdefault("corrections", {})
                step_corrs.setdefault(step_name, []).extend(step_corrections)

        # 구분자 원장으로 접합 -- 절단면의 공백/개행을 그대로 복원한다. 분할·접합 규칙은 chunking.py.
        ctx.text = join_chunks(processed_bulks, plan)

    # ── helpers (multi_step.py 에서 그대로 이동) ─────────────────────────

    def _extract_output(self, content: str) -> str:
        """Extract output from LLM response. Tries JSON parsing first.

        Handles cases where model outputs reasoning/steps before the final JSON.
        Searches for the last valid JSON object containing an 'output' key.
        """
        # Try raw JSON (entire content is valid JSON)
        try:
            data = json.loads(content)
            if isinstance(data, dict) and "output" in data:
                return data["output"]
        except (json.JSONDecodeError, AttributeError):
            pass

        # Try markdown code blocks — find the last one with 'output' key
        json_blocks = list(re.finditer(r"```(?:json)?\s*(.*?)```", content, re.DOTALL))
        for match in reversed(json_blocks):
            try:
                data = json.loads(match.group(1).strip())
                if isinstance(data, dict) and "output" in data:
                    return data["output"]
            except (json.JSONDecodeError, AttributeError):
                continue

        # Try finding JSON by scanning for last {"output": pattern
        # Avoid complex regex to prevent catastrophic backtracking on long texts
        last_idx = content.rfind('"output"')
        if last_idx != -1:
            # Find the enclosing { before "output"
            brace_start = content.rfind("{", 0, last_idx)
            if brace_start != -1:
                # Find matching } by counting braces
                depth = 0
                for i in range(brace_start, len(content)):
                    if content[i] == "{":
                        depth += 1
                    elif content[i] == "}":
                        depth -= 1
                        if depth == 0:
                            try:
                                data = json.loads(content[brace_start : i + 1])
                                if isinstance(data, dict) and "output" in data:
                                    return data["output"]
                            except (json.JSONDecodeError, AttributeError):
                                pass
                            break

        # Try extracting "output" from truncated JSON
        # When explanation is too long, JSON gets truncated by max_tokens,
        # but the output field (which comes first) is usually intact.
        # Pattern: "output": "...value...", "explanation": ...
        output_key_match = re.search(r'"output"\s*:\s*"', content)
        if output_key_match:
            start = output_key_match.end()
            remaining = content[start:]
            # Find boundary: closing quote of output value followed by "explanation" key
            boundary = re.search(r'"\s*,\s*"explanation"\s*:', remaining)
            if boundary:
                raw_value = remaining[: boundary.start()]
                try:
                    return json.loads('"' + raw_value + '"')
                except (json.JSONDecodeError, ValueError):
                    return raw_value

        return content

    async def _run_self_consistency_bulk(
        self,
        bulk: str,
        step: dict[str, Any],
        *,
        step_model: str,
        step_re: str | None,
        step_max_tokens: int,
        step_response_format: dict[str, Any] | None,
        step_prompt: str,
        user_template: str | None,
        include_original: bool,
        original_text: str,
        step_outputs: dict[str, str],
        ctx: StepContext,
    ) -> tuple[str, dict[str, int], list[dict[str, Any]], list[dict[str, Any]]]:
        """basic_correction 스텝의 self-consistency 실행. 프롬프트 조립만 여기서 하고
        반복 호출/교집합/적용은 `judge.run_self_consistency` 에 위임한다 (§5-F).
        """
        if user_template:
            user_content = self._format_user_template(
                user_template, original_text, bulk, step_outputs
            )
        elif include_original:
            user_content = f"[원문]\n{original_text}\n\n[교정본]\n{bulk}"
        else:
            user_content = bulk

        return await judge.run_self_consistency(
            bulk,
            user_content,
            step,
            provider=ctx.provider,
            step_model=step_model,
            step_re=step_re,
            step_max_tokens=step_max_tokens,
            step_response_format=step_response_format,
            step_prompt=step_prompt,
        )

    def _format_user_template(
        self,
        template: str,
        original: str,
        current: str,
        step_outputs: dict[str, str],
    ) -> str:
        """Format user content template with step outputs.

        Supported placeholders:
            {original}       — raw original text from dataset
            {current}        — current pipeline text (output of previous step)
            {steps.<name>}   — output of a specific named step
        """
        result = template.replace("{original}", original).replace("{current}", current)
        for name, output in step_outputs.items():
            result = result.replace(f"{{steps.{name}}}", output)
        return result