# app.py
from typing import List, Dict, Any, Optional
import re
import json

import gradio as gr
from transformers import pipeline

# =========================
# モデル定義（Hugging Face）
# =========================

# 日本語感情分析モデル
SENTIMENT_MODEL_NAME = "koheiduck/bert-japanese-finetuned-sentiment"
# 多言語ゼロショット分類モデル（STAR分類・トーン評価用）
NLI_MODEL_NAME = "joeddav/xlm-roberta-large-xnli"

# 必要に応じて device=0 (GPU) に変更してください
sentiment_classifier = pipeline(
    "sentiment-analysis",
    model=SENTIMENT_MODEL_NAME,
    tokenizer=SENTIMENT_MODEL_NAME,
)

zero_shot_classifier = pipeline(
    "zero-shot-classification",
    model=NLI_MODEL_NAME,
    tokenizer=NLI_MODEL_NAME,
)

# =========================
# ユーティリティ
# =========================

_SENT_SPLIT_RE = re.compile(r"[。！？\n]+")


def split_sentences(text: str) -> List[str]:
    """日本語の文をざっくり分割する簡易関数"""
    sentences = [s.strip() for s in _SENT_SPLIT_RE.split(text) if s.strip()]
    return sentences


# =========================
# 感情・トーン分析
# =========================

def analyze_sentiment(text: str) -> Dict[str, Any]:
    """
    日本語テキストのネガポジ分析。
    出力例:
    {
        "label": "POSITIVE",
        "score": 0.98
    }
    """
    if not text.strip():
        return {"label": "NEUTRAL", "score": 0.0}

    result = sentiment_classifier(text)[0]
    return {
        "label": result["label"],
        "score": float(result["score"]),
    }


def analyze_tone(text: str) -> Dict[str, Any]:
    """
    ゼロショット分類で「熱意・主体性・一貫性・論理性・協調性」などをスコアリング。
    multi_label=True なので、1文が複数のトーンを兼ねることも許容。
    出力例:
    {
        "labels": {
            "熱意": 0.92,
            "主体性": 0.81,
            "一貫性": 0.55,
            "論理性": 0.73,
            "協調性": 0.40
        }
    }
    """
    if not text.strip():
        return {"labels": {}}

    tone_labels = ["熱意", "主体性", "一貫性", "論理性", "協調性"]

    result = zero_shot_classifier(
        text,
        candidate_labels=tone_labels,
        multi_label=True,
    )

    label_scores = {
        label: float(score)
        for label, score in zip(result["labels"], result["scores"])
    }

    ordered = {label: label_scores.get(label, 0.0) for label in tone_labels}

    return {"labels": ordered}


# =========================
# STAR 構造分析
# =========================

STAR_LABELS_JP = [
    "Situation(状況)",
    "Task(課題)",
    "Action(行動)",
    "Result(結果)",
    "Other(その他)",
]

STAR_KEY_MAP = {
    "Situation(状況)": "S",
    "Task(課題)": "T",
    "Action(行動)": "A",
    "Result(結果)": "R",
    "Other(その他)": "O",
}


def classify_sentence_star(sentence: str) -> Dict[str, Any]:
    """
    1文を STAR のどれに近いかゼロショット分類する。
    multi_label=False とし、最も近いラベルのみ採用。
    出力例:
    {
        "sentence": "...",
        "star_label": "S",
        "raw_label": "Situation(状況)",
        "score": 0.87
    }
    """
    if not sentence.strip():
        return {
            "sentence": sentence,
            "star_label": "O",
            "raw_label": "Other(その他)",
            "score": 0.0,
        }

    result = zero_shot_classifier(
        sentence,
        candidate_labels=STAR_LABELS_JP,
        multi_label=False,
    )
    raw_label = result["labels"][0]
    score = float(result["scores"][0])
    star_label = STAR_KEY_MAP.get(raw_label, "O")

    return {
        "sentence": sentence,
        "star_label": star_label,
        "raw_label": raw_label,
        "score": score,
    }


def analyze_star_structure(text: str) -> Dict[str, Any]:
    """
    自己PRテキストの STAR 構造（S/T/A/R がどの程度含まれているか）を分析。
    出力例:
    {
        "coverage": {
            "S": true,
            "T": true,
            "A": true,
            "R": false
        },
        "star_score": 0.75,
        "per_sentence": [...],
        "missing_elements": ["R"],
        "comment": "結果(Result)の記述が弱いため、成果をより具体的に書くと良いです。"
    }
    """
    sentences = split_sentences(text)
    if not sentences:
        return {
            "coverage": {k: False for k in ["S", "T", "A", "R"]},
            "star_score": 0.0,
            "per_sentence": [],
            "missing_elements": ["S", "T", "A", "R"],
            "comment": "文章が空か極端に短いため、STAR 構造の判定ができません。",
        }

    per_sentence_results = [
        classify_sentence_star(s) for s in sentences
    ]

    coverage = {k: False for k in ["S", "T", "A", "R"]}
    for r in per_sentence_results:
        key = r["star_label"]
        if key in coverage:
            coverage[key] = True

    missing = [k for k, v in coverage.items() if not v]
    star_score = sum(1 for v in coverage.values() if v) / 4.0

    # コメント生成（シンプルなヒューリスティック）
    if not missing:
        comment = "STAR の各要素（状況・課題・行動・結果）が一通り含まれています。構成としてバランスは良好です。"
    else:
        parts = []
        mapping_jp = {"S": "状況(Situation)", "T": "課題(Task)", "A": "行動(Action)", "R": "結果(Result)"}
        for m in missing:
            parts.append(mapping_jp[m])
        missing_jp = "・".join(parts)
        comment = (
            f"{missing_jp} の要素が弱い／不足しています。"
            "不足している要素を具体的に書き足すと、より論理的で説得力のある自己PRになります。"
        )

    return {
        "coverage": coverage,
        "star_score": star_score,
        "per_sentence": per_sentence_results,
        "missing_elements": missing,
        "comment": comment,
    }


# =========================
# ES 全体の評価インターフェース
# =========================

def evaluate_entry_sheet(
    self_pr: str,
    motivation: Optional[str] = None,
) -> Dict[str, Any]:
    """
    ES の自己PR（必須）、志望動機（任意）を入力として、
    感情・トーン・STAR 構造をまとめて評価するインターフェース。
    """
    texts = [self_pr]
    if motivation:
        texts.append(motivation)
    full_text = "\n".join([t for t in texts if t])

    sentiment = analyze_sentiment(full_text)
    tone = analyze_tone(full_text)
    star = analyze_star_structure(self_pr)  # STAR は主に自己PRに対して実行

    return {
        "input": {
            "self_pr": self_pr,
            "motivation": motivation,
        },
        "sentiment": sentiment,
        "tone": tone,
        "star": star,
    }


# =========================
# Gradio 用ラッパー
# =========================

def evaluate_es(self_pr: str, motivation: str):
    """
    Gradio から呼ばれるラッパー関数。
    self_pr: 自己PR
    motivation: 志望動機（空でもOK）
    """
    if not self_pr.strip() and not motivation.strip():
        return {
            "error": "自己PRか志望動機のいずれかは入力してください。"
        }

    result = evaluate_entry_sheet(
        self_pr=self_pr,
        motivation=motivation if motivation.strip() else None,
    )
    return result


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 新卒エントリーシート AI評価デモ（Hugging Face 上で完結）

        - 日本語感情分析モデル + 多言語ゼロショット分類モデルを用いて、
          ES の **感情・トーン** と **STAR 構造（Situation / Task / Action / Result）** を自動評価します。
        - 出力は JSON 形式なので、そのまま既存パイプラインに組み込み可能です。
        """
    )

    with gr.Row():
        self_pr_input = gr.Textbox(
            label="自己PR（必須）",
            lines=12,
            placeholder="自己PRを入力してください。",
        )
        motivation_input = gr.Textbox(
            label="志望動機（任意）",
            lines=12,
            placeholder="志望動機があれば入力してください。（空でも可）",
        )

    btn = gr.Button("AIで評価する")

    output_json = gr.JSON(
        label="評価結果（JSON）",
        value={},
    )

    btn.click(
        fn=evaluate_es,
        inputs=[self_pr_input, motivation_input],
        outputs=output_json,
    )

if __name__ == "__main__":
    demo.launch()