# -*- coding: utf-8 -*-
"""Step 5: LLM-assisted annotation (OpenAI Batch API)

입력:
  - output/{corpus_type}/step4_candidates_with_examples.csv
출력:
  - output/{corpus_type}/step5_llm_assisted_annotation.csv
  - output/{corpus_type}/step5_llm_assisted_annotation/
    ├── review.xlsx
    ├── log.json
    ├── batch_input.jsonl
    ├── batch_output.jsonl
    └── batch_status.json
"""

from __future__ import annotations

import json
import os
import time
from pathlib import Path

import pandas as pd
from openpyxl.styles import PatternFill, Font, Alignment

from wordextractor.config import PipelineConfig
from wordextractor.prompt import format_prompt, parse_judgment_json, RESULT_KEYS, _EMPTY_RESULT


def _get_openai_client(cfg: PipelineConfig):
    try:
        from openai import OpenAI
    except ImportError:
        raise ImportError(
            "Step 5 requires the 'llm' extra. "
            "Install with: pip install wordextractor[llm]"
        )

    if cfg.openai_api_key:
        os.environ.setdefault("OPENAI_API_KEY", cfg.openai_api_key)
    elif cfg.openai_env_path:
        try:
            from dotenv import load_dotenv

            load_dotenv(cfg.openai_env_path)
        except ImportError:
            raise ImportError(
                "Step 5 requires python-dotenv for .env file loading. "
                "Install with: pip install wordextractor[llm]"
            )

    return OpenAI()


def build_batch_input(df: pd.DataFrame, input_path: Path, model: str, temperature: float) -> None:
    with open(input_path, "w", encoding="utf-8") as f:
        for idx, row in df.iterrows():
            word = row["어절"]
            morphs = str(row.get("형태소 분석", "") or "")
            pattern = str(row.get("매칭 패턴", "") or "")
            exs = [
                str(row.get("용례1", "") or ""),
                str(row.get("용례2", "") or ""),
                str(row.get("용례3", "") or ""),
            ]
            prompt = format_prompt(word, morphs, pattern, exs)

            request = {
                "custom_id": str(idx),
                "method": "POST",
                "url": "/v1/chat/completions",
                "body": {
                    "model": model,
                    "messages": [{"role": "user", "content": prompt}],
                    "temperature": temperature,
                },
            }
            f.write(json.dumps(request, ensure_ascii=False) + "\n")

    print(f"  Batch 입력 저장: {input_path} ({len(df):,}건)")


def submit_batch(client, input_path: Path, status_path: Path) -> str:
    print("  파일 업로드 중...")
    with open(input_path, "rb") as f:
        file_obj = client.files.create(file=f, purpose="batch")
    print(f"  파일 ID: {file_obj.id}")

    print("  Batch 생성 중...")
    batch = client.batches.create(
        input_file_id=file_obj.id,
        endpoint="/v1/chat/completions",
        completion_window="24h",
    )
    print(f"  Batch ID: {batch.id}")

    _save_status(status_path, {"batch_id": batch.id, "input_file_id": file_obj.id, "status": batch.status})
    return batch.id


def wait_for_batch(client, batch_id: str, status_path: Path) -> dict:
    poll_interval = 30
    while True:
        batch = client.batches.retrieve(batch_id)
        status = batch.status
        completed = batch.request_counts.completed if batch.request_counts else 0
        total = batch.request_counts.total if batch.request_counts else 0

        print(f"  상태: {status} ({completed}/{total})", end="\r")

        _save_status(status_path, {
            "batch_id": batch_id,
            "status": status,
            "completed": completed,
            "total": total,
        })

        if status in ("completed", "failed", "cancelled", "expired"):
            print(f"\n  최종 상태: {status}")
            return {
                "status": status,
                "output_file_id": batch.output_file_id,
                "error_file_id": batch.error_file_id,
            }

        time.sleep(poll_interval)


def download_output(client, output_file_id: str, output_path: Path) -> None:
    content = client.files.content(output_file_id)
    with open(output_path, "wb") as f:
        f.write(content.read())
    print(f"  Batch 출력 저장: {output_path}")


def parse_batch_output(df: pd.DataFrame, output_path: Path) -> pd.DataFrame:
    results = {}
    with open(output_path, "r", encoding="utf-8") as f:
        for line in f:
            obj = json.loads(line)
            custom_id = int(obj["custom_id"])
            response = obj.get("response", {})
            body = response.get("body", {})
            choices = body.get("choices", [])

            if choices:
                text = choices[0].get("message", {}).get("content", "")
                result = parse_judgment_json(text)
            else:
                result = dict(_EMPTY_RESULT)
                error = obj.get("error", {})
                result["판정 근거"] = f"API 오류: {error.get('message', 'unknown')}"

            results[custom_id] = result

    for key in RESULT_KEYS:
        df[key] = df.index.map(lambda i, k=key: results.get(i, _EMPTY_RESULT).get(k, ""))

    return df


def save_review_xlsx(df: pd.DataFrame, review_path: Path) -> None:
    df_review = df[
        [
            "어절",
            "매칭 패턴",
            "형태소 분석",
            "판정 결과",
            "판정 기준",
            "판정 근거",
            "혼성어_어형",
            "원어X",
            "원어X_원어",
            "원어X_어종",
            "원어Y",
            "원어Y_원어",
            "원어Y_어종",
            "용례1",
            "용례2",
            "용례3",
        ]
    ].copy()

    df_review.insert(0, "idx", range(len(df_review)))
    df_review.insert(2, "최종 판단", "")
    df_review.insert(3, "메모", "")

    with pd.ExcelWriter(review_path, engine="openpyxl") as writer:
        df_review.to_excel(writer, index=False, sheet_name="Sheet1")

        ws = writer.sheets["Sheet1"]
        ws.freeze_panes = "A2"
        ws.auto_filter.ref = ws.dimensions

        header_fill = PatternFill(start_color="B4C6E7", end_color="B4C6E7", fill_type="solid")
        header_font = Font(bold=True)

        for cell in ws[1]:
            cell.fill = header_fill
            cell.font = header_font
            cell.alignment = Alignment(horizontal="center", vertical="center")

    print(f"리뷰용 결과 저장: {review_path}")


def save_log(df: pd.DataFrame, log_path: Path, model: str, temperature: float) -> dict:
    n_blend = len(df[df["판정 결과"] == "혼성어"])
    n_non = len(df[df["판정 결과"] == "비혼성어"])
    n_unk = len(df[df["판정 결과"] == "판정 불가"])
    total = len(df)

    summary = {
        "judgment_model": model,
        "temperature": temperature,
        "total": total,
        "혼성어": n_blend,
        "혼성어 비율": round(n_blend / total * 100, 2) if total else 0,
        "비혼성어": n_non,
        "비혼성어 비율": round(n_non / total * 100, 2) if total else 0,
        "판정 불가": n_unk,
        "판정 불가 비율": round(n_unk / total * 100, 2) if total else 0,
    }

    with open(log_path, "w", encoding="utf-8") as f:
        json.dump(summary, f, ensure_ascii=False, indent=2)
    print(f"로그 저장: {log_path}")

    return summary


def _save_status(status_path: Path, data: dict) -> None:
    with open(status_path, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def _load_status(status_path: Path) -> dict:
    if status_path.exists():
        with open(status_path, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}


def run(cfg: PipelineConfig) -> None:
    out_dir = cfg.ensure_output_dir()
    step5_dir = out_dir / "step5_llm_assisted_annotation"
    step5_dir.mkdir(parents=True, exist_ok=True)

    input_path = step5_dir / "batch_input.jsonl"
    output_path = step5_dir / "batch_output.jsonl"
    status_path = step5_dir / "batch_status.json"
    review_path = step5_dir / "review.xlsx"
    log_path = step5_dir / "log.json"
    result_path = out_dir / "step5_llm_assisted_annotation.csv"

    client = _get_openai_client(cfg)

    print("데이터 로드...")
    df = pd.read_csv(out_dir / "step4_candidates_with_examples.csv", encoding="utf-8-sig")
    print(f"  후보: {len(df):,}개")

    # 이전 배치 상태 확인
    status = _load_status(status_path)

    if status.get("status") == "completed" and output_path.exists():
        print(f"\n이전 Batch 완료 상태 발견. 결과 파싱 중...")
    else:
        batch_id = status.get("batch_id")

        if batch_id and status.get("status") not in ("completed", "failed", "cancelled", "expired"):
            print(f"\n진행 중인 Batch 발견: {batch_id}")
        else:
            print("\nBatch 입력 생성...")
            build_batch_input(df, input_path, cfg.judgment_model, cfg.temperature)

            print("\nBatch 제출...")
            batch_id = submit_batch(client, input_path, status_path)

        print("\nBatch 완료 대기...")
        result = wait_for_batch(client, batch_id, status_path)

        if result["status"] != "completed":
            print(f"Batch 실패: {result['status']}")
            if result.get("error_file_id"):
                err_content = client.files.content(result["error_file_id"])
                err_path = out_dir / "step5_batch_errors.jsonl"
                with open(err_path, "wb") as f:
                    f.write(err_content.read())
                print(f"  오류 파일: {err_path}")
            return

        print("\nBatch 결과 다운로드...")
        download_output(client, result["output_file_id"], output_path)

    # 결과 파싱 및 병합
    print("\n결과 파싱...")
    df = parse_batch_output(df, output_path)

    # CSV 저장
    df.to_csv(result_path, index=False, encoding="utf-8-sig")
    print(f"결과 저장: {result_path}")

    # 리뷰용 XLSX 저장
    save_review_xlsx(df, review_path)

    # 로그 JSON 저장
    summary = save_log(df, log_path, cfg.judgment_model, cfg.temperature)

    # 요약 출력
    print(f"\n=== 요약 ===")
    print(f"  전체: {summary['total']:,}개")
    print(f"  혼성어: {summary['혼성어']:,}개 ({summary['혼성어 비율']}%)")
    print(f"  비혼성어: {summary['비혼성어']:,}개 ({summary['비혼성어 비율']}%)")
    print(f"  판정 불가: {summary['판정 불가']:,}개 ({summary['판정 불가 비율']}%)")
