# -*- coding: utf-8 -*-
"""
会議TXTを読み込み → 要約プロンプトを自己評価で最適化（教師なし）
評価軸:
  1) 簡潔さ (conciseness)
  2) 事実の包含量 (coverage_of_facts)
  3) 決定事項の網羅 (decisions_completeness)
  4) 誤情報の少なさ (factuality / hallucination avoidance)

要件:
- ./meetings/*.txt を読み込む。1ファイルしか無ければ自動チャンク分割。
- LLM-as-a-Judge（同じLMでも別LMでもOK）で4軸を各0〜5点のJSONで採点し、総合スコアをメトリクとして最適化。
- 出力は最適化前後のスコア比較、代表サンプルの要約、最適化後プロンプト表示。

注:
- “judge”もAPIコールなので、コスト節約したい場合は judge を軽いモデルにする/評価件数を減らす/要約長を短くする等で調整してください。
"""

import os, glob, math, random, re, json, textwrap
from typing import List, Dict
import dspy

# ========= LLM 設定 =========
# 例: OpenAI (LiteLLM互換) / Ollama / Azure OpenAI など
# judge と generator を分けるのが理想（同じでもOK）
GEN_MODEL = "openai/gpt-5-mini"   # 要約を作るLM
JUDGE_MODEL = "openai/gpt-4o-mini" # 自己評価するLM（別にできるなら別を推奨）

gen_lm = dspy.LM(GEN_MODEL, temperature=1.0, max_tokens=16000)
judge_lm = dspy.LM(JUDGE_MODEL)
dspy.configure(lm=gen_lm)

# ========= データ読み込み =========
DATA_DIR = "./meetings"
MAX_CHARS_PER_DOC = 4000   # 1サンプルの入力長上限（長すぎる会議は分割）
OVERLAP = 400              # 分割時のオーバーラップ

def load_or_chunk_texts() -> List[Dict]:
    paths = sorted(glob.glob(os.path.join(DATA_DIR, "*.txt")))
    if not paths:
        raise FileNotFoundError("meetings/*.txt が見つかりません。会議テキストを置いてください。")
    samples = []
    for p in paths:
        txt = open(p, "r", encoding="utf-8").read().strip()
        if not txt: continue
        if len(txt) <= MAX_CHARS_PER_DOC:
            samples.append({"doc_id": os.path.basename(p), "text": txt})
        else:
            # ざっくり文字数でチャンク分割（段落ベースにしたいならここを改良）
            i = 0; n = 1
            while i < len(txt):
                chunk = txt[i:i+MAX_CHARS_PER_DOC]
                samples.append({"doc_id": f"{os.path.basename(p)}#part{n}", "text": chunk})
                i += MAX_CHARS_PER_DOC - OVERLAP
                n += 1
    # シャッフルして返す
    random.shuffle(samples)
    return samples

records = load_or_chunk_texts()

# ========= DSPy の Signature / Module =========
class MeetingSummSig(dspy.Signature):
    """あなたは優れた会議要約アシスタントです。
    入力の議事録テキストを読み、以下の形式で日本語の要約を作成してください。

    1) 概要（1～3文で会議の全体像）
    2) 決定事項（箇条書き、漏れなく。なければ「なし」と明記）
    3) アクションアイテム（担当者/期日を含め箇条書き。なければ「なし」と明記）

    制約:
    - 事実のみ。推測や脚色は禁止。
    - 固有名や数値は可能な限り保持。
    - 全体で日本語 200～400文字程度（調整可）。
    """
    transcript: str = dspy.InputField()
    summary: str = dspy.OutputField(desc="日本語の構造化要約。上記フォーマットを守ること。")

class MeetingSummarizer(dspy.Module):
    def __init__(self):
        super().__init__()
        self.cot = dspy.ChainOfThought(MeetingSummSig)

    def forward(self, transcript: str):
        pred = self.cot(transcript=transcript[:MAX_CHARS_PER_DOC])
        # 軽いクレンジング（全角/半角や空白の体裁調整など、必要なら追加）
        out = re.sub(r"\n{3,}", "\n\n", pred.summary).strip()
        return dspy.Prediction(summary=out)

summarizer = MeetingSummarizer()

# ========= 自己評価（LLM-as-a-Judge） =========
JUDGE_PROMPT = """あなたは厳密な評価者です。元の議事録（SOURCE）と要約（SUMMARY）を読み、次の4軸で0〜5点の整数で採点してください。
- conciseness: 簡潔さ（冗長性が少なく要点がまとまっている）
- coverage_of_facts: 事実の包含量（重要な事実・数値・固有名が適切に含まれている）
- decisions_completeness: 決定事項の網羅（決定内容が漏れなく明示されている）
- factuality: 事実性（SOURCEにない断定や誤りが無い、矛盾しない）

必ず次のJSONだけを出力してください（説明文なし）:
{{"conciseness": <0~5>, "coverage_of_facts": <0~5>, "decisions_completeness": <0~5>, "factuality": <0~5>}}

SOURCE:
<<<
{source}
>>>

SUMMARY:
<<<
{summary}
>>>
"""


def judge_scores(source: str, summary: str) -> Dict[str, int]:
    # LLMに厳密JSONを要求
    prompt = JUDGE_PROMPT.format(source=source[:MAX_CHARS_PER_DOC], summary=summary)
    raw = judge_lm(prompt)
    txt = str(raw)
    # JSON抽出（多少のノイズに耐える）
    m = re.search(r'\{.*\}', txt, flags=re.DOTALL)
    js = {"conciseness":0,"coverage_of_facts":0,"decisions_completeness":0,"factuality":0}
    if m:
        try:
            js = json.loads(m.group(0))
        except Exception:
            pass
    # ガード: 範囲外はクリップし整数化
    for k in js.keys():
        try:
            v = int(round(float(js[k])))
        except Exception:
            v = 0
        js[k] = max(0, min(5, v))
    return js

# 総合スコア（重みは必要に応じて調整）
WEIGHTS = {
    "conciseness": 0.2,
    "coverage_of_facts": 0.3,
    "decisions_completeness": 0.3,
    "factuality": 0.2,
}
def aggregate_score(js: Dict[str,int]) -> float:
    # 0〜1に正規化
    s = sum(WEIGHTS[k] * (js[k]/5.0) for k in WEIGHTS)
    return float(s)

# ========= データをExamplesへ =========
def to_example(item: Dict) -> dspy.Example:
    return dspy.Example(
        transcript=item["text"],
    ).with_inputs("transcript")

# 学習/開発/テスト 分割（教師なしだが汎化確認用に分ける）
def split(items, tr=0.6, dev=0.2, seed=42):
    random.Random(seed).shuffle(items)
    n = len(items); n_tr = int(n*tr); n_dev = int(n*dev)
    return items[:n_tr], items[n_tr:n_tr+n_dev], items[n_tr+n_dev:]

train, dev, test = split(records)
train_ex = [to_example(r) for r in train]
dev_ex   = [to_example(r) for r in dev]

# ========= メトリクス（自己評価） =========
try:
    from dspy.teleprompt.gepa.gepa import ScoreWithFeedback
except Exception:
    ScoreWithFeedback = None

def gepa_metric(gold: dspy.Example,
                pred: dspy.Prediction,
                trace=None,
                pred_name: str | None = None,
                pred_trace=None):
    # 入力抽出（堅牢）
    src = getattr(gold, "transcript", None)
    if not src:
        ins_attr = getattr(gold, "inputs", None)
        ins = ins_attr() if callable(ins_attr) else ins_attr
        if isinstance(ins, dict):
            src = ins.get("transcript")
    if not src:
        return 0.0

    summary = pred.summary
    js = judge_scores(src, summary)  # {"conciseness":..., ...}
    score = aggregate_score(js)

    # 予測器名が来てたら、それにひもづく短いFBを作る
    fb_lines = [
        f"[{pred_name or 'program'}] score={score:.3f}",
        f"- conciseness: {js['conciseness']}/5",
        f"- coverage_of_facts: {js['coverage_of_facts']}/5",
        f"- decisions_completeness: {js['decisions_completeness']}/5",
        f"- factuality: {js['factuality']}/5",
    ]
    # 改善ヒント（例）
    if js["decisions_completeness"] < 4:
        fb_lines.append("→ 決定事項を『箇条書きで漏れなく・固有名と期日込み』で明示して。")
    if js["coverage_of_facts"] < 4:
        fb_lines.append("→ 固有名・数値の保持を優先（曖昧表現を避ける）。")
    if js["factuality"] < 5:
        fb_lines.append("→ SOURCEに無い断定や推測は禁止。文章を該当箇所に絞る。")
    if js["conciseness"] < 4:
        fb_lines.append("→ 200〜400字目安を厳守し、重複表現を圧縮。")

    feedback = "\n".join(fb_lines)

    if ScoreWithFeedback is not None:
        return ScoreWithFeedback(score=score, feedback=feedback)
    else:
        return score


# ========= ベースライン評価 =========
def eval_set(items: List[Dict], program: MeetingSummarizer, k_show=2, tag=""):
    total = 0.0
    shown = 0
    for it in items:
        out = program(transcript=it["text"])
        js = judge_scores(it["text"], out.summary)
        sc = aggregate_score(js)
        total += sc
        if shown < k_show:
            print(f"\n[{tag} sample] {it['doc_id']}")
            print("scores:", js, "| aggregate:", round(sc, 3))
            print("summary:\n", out.summary[:800])
            shown += 1
    avg = total / max(1, len(items))
    return avg

print("=== Baseline ===")
baseline_dev  = eval_set(dev, summarizer, tag="dev baseline")
baseline_test = eval_set(test, summarizer, tag="test baseline")
print(f"\nDev avg score (baseline):  {baseline_dev:.3f}")
print(f"Test avg score (baseline): {baseline_test:.3f}")

# ==== GEPA オプティマイザに差し替え ====
from dspy.teleprompt.gepa import GEPA   # ← モジュール名はこれで入るはず
gepa = GEPA(
    metric=gepa_metric,
    auto="medium",           # ← おまかせにするなら指定。手動回数を使うなら auto=None
    reflection_minibatch_size=2,
    reflection_lm=judge_lm
)
summarizer_opt = gepa.compile(
    summarizer,               # ← あなたの dspy.Module
    trainset=train_ex,        # ← キーワード引数で
    valset=dev_ex,            # ← 任意
    # num_trials=12,          # auto を使うなら外す。手動なら auto=None にして指定
)


print("\n--- Optimized Prompt (excerpt) ---")
try:
    summarizer_opt.save(f"optimized.json")
    params = summarizer_opt.get_params()
    print(json.dumps(params, ensure_ascii=False, indent=2))
except Exception:
    pass

print("\n=== After Optimization ===")
opt_dev  = eval_set(dev, summarizer_opt, tag="dev opt")
opt_test = eval_set(test, summarizer_opt, tag="test opt")
print(f"\nDev avg score (opt):  {opt_dev:.3f}")
print(f"Test avg score (opt): {opt_test:.3f}")


# ========= 推論用API（サービス統合の最小形） =========
def summarize_meeting(text: str) -> Dict[str, str]:
    """本番システムから呼ぶ用：最適化後の要約器を1関数でラップ"""
    out = summarizer_opt(transcript=text)
    return {
        "summary": out.summary
    }

# 例:
# ex = summarize_meeting(open("./meetings/example.txt","r",encoding="utf-8").read())
# print(ex["summary"])