"""파이프라인 = 등록된 스텝들을 순차 실행하는 실행기 **하나**.

예전에는 SingleStepPipeline/MultiStepPipeline 두 클래스가 각자 실행 루프를
가졌고, 스텝 종류는 그 안의 if/elif 분기였다. 이제 실행 루프는 이 파일의
`Pipeline` 하나뿐이고, 태스크별 차이는 전부 "어떤 스텝을 어떤 순서로
꽂는가"(steps 목록 × STEP_REGISTRY)로 표현된다. v24 처럼 LLM 스텝 사이사이
rule 후처리 블록이 끼어드는 구조가 이 모델의 원형이다.

스텝의 계약:
  - `ctx.text` 가 체인의 현재 텍스트다. 스텝은 이걸 읽고 덮어쓴다.
  - 구조화 부산물(교정 목록, judge 결정 등)은 `sample.artifacts` 에 남긴다.
  - LLM 사용량은 `ctx.add_usage()` 로 합산한다.
  - 실행기가 스텝마다 `step_outputs[스텝이름] = ctx.text` 스냅숏을 찍는다.
"""

import asyncio
import logging
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from pathlib import Path
from collections.abc import Callable
from typing import Any

from solar_eval.core.dataset_loader import DatasetLoader
from solar_eval.models.sample import EvalSample
from solar_eval.pipelines.base import BasePipeline
from solar_eval.providers.base import BaseProvider

logger = logging.getLogger(__name__)


@dataclass
class PipelineEnv:
    """파이프라인 생성 시점에 고정되는 정적 환경. 스텝 생성자에 전달된다."""

    prompts: dict[str, str] = field(default_factory=dict)
    input_fields: list[str] = field(default_factory=list)
    text_split_size: int = 500
    # bulk 분할 전략 -- "dp"(운영) | "greedy"(regex 사다리, ablation 전용). steps/chunking.py 참조.
    text_split_strategy: str = "dp"
    # 문서 단위 chunk (실험, 2026-09-10): 입구에서 한 번 DP 분할해 chunk 마다 **전 스텝**을 돌리고 원장으로 접합.
    # 켜면 분할 없이 전문을 받던 스텝(basic_correction 의 user_template)도 chunk 의 원문만 본다.
    document_chunking: bool = False
    # 문서 단위 chunk 를 동시에 몇 개까지 돌릴지. chunk 마다 LLM 호출이 순차로 이어지므로 이 값 × 러너 워커 수가
    # 동시 API 호출 상한이다 (solar-pro4 는 10 동시부터 429 가 났다 -- README "지뢰").
    document_chunk_concurrency: int = 4
    dataset_loader: DatasetLoader | None = None
    # 데이터 루트의 프로젝트 디렉토리. v24 의 tool_calling_judge 가 pmi_lookup
    # (1.5MB 파생 통계 — 사람이 편집하지 않는 데이터 자산) 을 풀 때 쓰인다.
    project_dir: Path | None = None
    # 레포의 config 정본 디렉토리(`accounts/<a>/<c>/03-evaluation`). 치환 사전처럼
    # **사람이 편집하는 지식 자산**은 데이터가 아니라 설정이라 여기서 읽는다.
    config_dir: Path | None = None


# 진행 콜백 -- (event, index, total, step_name, doc_chunk). 데모 UI 가 "지금 어느 단계인가" 를 보여주는 용도,
# 결과에는 아무 영향이 없다.
#   event: "start" | "done" (스텝), "chunk" (LLM 스텝의 bulk 진행), "document" (문서 단위 chunk 시작, step_name 은 "")
#   doc_chunk: 문서 단위 chunk 모드에서 (chunk 번호, chunk 수), 아니면 None. chunk 들이 **병렬**로 돌므로
#   이벤트가 섞여 들어온다 -- UI 는 이 값으로 줄을 나눠야 한다.
DocChunk = tuple[int, int] | None
StepCallback = Callable[[str, int, int, str, DocChunk], None]


@dataclass
class StepContext:
    """샘플 1건이 체인을 흐르는 동안의 가변 상태."""

    provider: BaseProvider
    model: str
    temperature: float
    max_tokens: int
    reasoning_effort: str | None
    # single_step 계열(run 의 prompts 인자)이 쓰던 시스템 프롬프트 문자열.
    system_prompt: str
    messages: list[dict[str, str]] | None
    original_text: str
    text: str
    step_outputs: dict[str, str] = field(default_factory=dict)
    # 파이프라인 수명 캐시 (vocabulary CSV 등) -- 샘플이 아니라 run 수명이다.
    shared: dict[Any, Any] = field(default_factory=dict)
    usage: dict[str, int] = field(
        default_factory=lambda: {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
    )
    on_step: StepCallback | None = None
    doc_chunk: DocChunk = None  # 문서 단위 chunk 모드에서 이 컨텍스트가 담당하는 chunk

    def add_usage(self, usage: dict[str, Any]) -> None:
        for k in self.usage:
            self.usage[k] += usage.get(k, 0)


class BaseStep(ABC):
    """등록 가능한 파이프라인 스텝.

    구현체는 `@register_step("이름")` 으로 STEP_REGISTRY 에 등록한다
    (`pipelines/registry.py`). YAML 의 step `type`(또는 legacy `rule_type`)이
    그 이름으로 조회되고, **파이프라인 생성 시점에** 전부 해석된다 --
    모르는 이름은 API 호출 한 번 없이 즉시 실패한다.
    """

    def __init__(self, config: dict[str, Any], env: PipelineEnv) -> None:
        self.config = config
        self.env = env
        self.name: str = config.get("name", config.get("type", type(self).__name__))

    @abstractmethod
    async def run(self, sample: EvalSample, ctx: StepContext) -> None:
        """`ctx.text` 를 변환하고/하거나 `sample.artifacts` 를 채운다."""
        ...


class Pipeline(BasePipeline):
    """스텝 체인 실행기. 파이프라인 종류 구분(single/multi)은 더 이상 없다."""

    def __init__(self, steps: list[BaseStep], env: PipelineEnv) -> None:
        self.steps = steps
        self.env = env
        # 스텝들이 run 수명으로 공유하는 캐시 (예: vocabulary CSV 1회 로드).
        self._shared: dict[Any, Any] = {}

    async def run(
        self,
        sample: EvalSample,
        *,
        prompts: str,
        provider: BaseProvider,
        model: str,
        temperature: float = 0.0,
        max_tokens: int = 8000,
        reasoning_effort: str | None = None,
        messages: list[dict[str, str]] | None = None,
        on_step: StepCallback | None = None,
    ) -> EvalSample:
        input_data = sample.input if isinstance(sample.input, dict) else {}
        original_text = input_data.get("original", "")
        start = time.monotonic()

        def make_ctx(text: str, doc_chunk: DocChunk = None) -> StepContext:
            return StepContext(
                provider=provider,
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
                reasoning_effort=reasoning_effort,
                system_prompt=prompts,
                messages=messages,
                original_text=text,
                text=text,
                shared=self._shared,
                on_step=on_step,
                doc_chunk=doc_chunk,
            )

        if self.env.document_chunking and len(original_text) > self.env.text_split_size:
            await self._run_by_document_chunks(sample, input_data, original_text, make_ctx, on_step)
        else:
            ctx = make_ctx(original_text)
            await self._run_chain(sample, ctx, on_step)
            sample.output = ctx.text
            sample.artifacts["usage"] = ctx.usage
            sample.artifacts["step_outputs"] = ctx.step_outputs

        sample.artifacts["inference_time_ms"] = (time.monotonic() - start) * 1000
        # completion 스텝은 자기 finish_reason 을 직접 남긴다 -- 없을 때만 기본값.
        sample.artifacts.setdefault("finish_reason", "stop")
        return sample

    async def _run_chain(
        self, sample: EvalSample, ctx: StepContext, on_step: StepCallback | None
    ) -> None:
        """스텝 목록을 순서대로 한 번 돈다. `ctx.text` 가 결과, `ctx.step_outputs` 가 스냅숏."""
        from datetime import datetime

        for step_i, step in enumerate(self.steps):
            ts = datetime.now().strftime("%H:%M:%S")
            logger.info(
                f"  [{ts}] step {step_i + 1}/{len(self.steps)} "
                f"'{step.name}' ({step.config.get('type', '?')}) start"
            )
            if on_step:
                on_step("start", step_i + 1, len(self.steps), step.name, ctx.doc_chunk)
            try:
                await step.run(sample, ctx)
            except Exception as e:
                ts = datetime.now().strftime("%H:%M:%S")
                logger.error(
                    f"  [{ts}] step {step_i + 1} '{step.name}' FAILED: {type(e).__name__}: {e}"
                )
                raise
            ctx.step_outputs[step.name] = ctx.text
            ts = datetime.now().strftime("%H:%M:%S")
            logger.info(f"  [{ts}] step {step_i + 1} '{step.name}' done, text_len={len(ctx.text)}")
            if on_step:
                on_step("done", step_i + 1, len(self.steps), step.name, ctx.doc_chunk)

    async def _run_by_document_chunks(
        self,
        sample: EvalSample,
        input_data: dict[str, Any],
        original_text: str,
        make_ctx: Callable[[str, DocChunk], StepContext],
        on_step: StepCallback | None,
    ) -> None:
        """문서를 입구에서 한 번 나눠 chunk 마다 전 스텝을 돌리고, 출력·스냅숏·사용량을 원장으로 합친다."""
        # steps/__init__ 가 llm.py 를 거쳐 이 모듈을 import 하므로 순환을 피해 지연 import
        from solar_eval.pipelines.steps.chunking import join_chunks, plan_chunks

        plan = plan_chunks(
            original_text, self.env.text_split_size, strategy=self.env.text_split_strategy
        )
        n_chunks = len(plan.chunks)
        limit = asyncio.Semaphore(max(1, int(self.env.document_chunk_concurrency)))

        async def run_chunk(ci: int, chunk: str) -> EvalSample:
            # chunk 끼리는 독립이라 동시에 돈다 -- 캐시(self._shared)는 읽기 전용으로만 공유
            doc_chunk: DocChunk = (ci + 1, n_chunks)
            async with limit:
                if on_step:
                    on_step("document", ci + 1, n_chunks, "", doc_chunk)
                sub = EvalSample(input={**input_data, "original": chunk})
                ctx = make_ctx(chunk, doc_chunk)
                await self._run_chain(sub, ctx, on_step)
                sub.output = ctx.text
                sub.artifacts["usage"] = ctx.usage
                sub.artifacts["step_outputs"] = ctx.step_outputs
                return sub

        subs: list[EvalSample] = list(
            await asyncio.gather(*(run_chunk(ci, c) for ci, c in enumerate(plan.chunks)))
        )

        sample.output = join_chunks([sub.output for sub in subs], plan)
        sample.artifacts["step_outputs"] = {
            name: join_chunks([sub.artifacts["step_outputs"][name] for sub in subs], plan)
            for name in subs[0].artifacts["step_outputs"]
        }
        usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
        for sub in subs:
            for k in usage:
                usage[k] += sub.artifacts["usage"].get(k, 0)
        sample.artifacts["usage"] = usage
        if any(sub.artifacts.get("finish_reason", "stop") != "stop" for sub in subs):
            sample.artifacts["finish_reason"] = "length"
        corrections: dict[str, list[Any]] = {}
        for sub in subs:
            for step_name, items in sub.artifacts.get("corrections", {}).items():
                corrections.setdefault(step_name, []).extend(items)
        if corrections:
            sample.artifacts["corrections"] = corrections
        sample.artifacts["document_chunks"] = {
            "lens": [len(c) for c in plan.chunks],
            "kinds": list(plan.kinds),
        }