import os
from typing import Any, Dict, List, Tuple

import gradio as gr
from openai import OpenAI


def get_openai_client() -> OpenAI:
    """OpenAI API クライアントを環境変数から初期化して返す。"""
    api_key = os.getenv("OPENAI_API_KEY", "").strip()
    if not api_key:
        raise RuntimeError("OPENAI_API_KEY が設定されていません。環境変数を確認してください。")

    base_url = os.getenv("OPENAI_API_BASE", "").strip() or None
    return OpenAI(api_key=api_key, base_url=base_url)


def get_vector_store_id() -> str:
    """環境変数 VECTOR_STORE_ID から単一のベクターストアIDを取得する。"""
    value = os.getenv("VECTOR_STORE_ID", "").strip()
    if not value:
        raise RuntimeError("VECTOR_STORE_ID が設定されていません。環境変数を確認してください。")
    return value.split(",")[0].strip()


def extract_citations(response_dict: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Responses API レスポンスから file_search の引用情報を抽出する。"""
    citations: List[Dict[str, Any]] = []
    outputs = response_dict.get("output") or []

    for output in outputs:
        for content in output.get("content", []) or []:
            annotations = content.get("annotations") or []
            for annotation in annotations:
                file_citation = annotation.get("file_citation")
                if not file_citation:
                    continue
                entry = {
                    "file_id": file_citation.get("file_id"),
                    "file_name": file_citation.get("file_name"),
                    "vector_store_id": file_citation.get("vector_store_id"),
                    "quote": file_citation.get("quote"),
                }
                if entry not in citations:
                    citations.append(entry)
    return citations


def answer_with_file_search(
    vector_store_id: str,
    query: str,
    top_k: int = 5,
    model: str = "gpt-4o-mini",
) -> Tuple[str, List[Dict[str, Any]]]:
    """
    Responses API + file_search ツールを用いて検索と回答を同時に行い、回答と引用を返す。

    Cookbook の "Integrating search results with LLM in a single API call" を参考に構成。
    """
    if not query.strip():
        return "", []

    client = get_openai_client()
    system_prompt = (
        "あなたは取得したドキュメントから回答するアシスタントです。"
    )

    user_message = f"{query}\n\n検索結果は最大 {top_k} 件まで引用してください。"

    response = client.responses.create(
        model=model,
        input=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message},
        ],
        tools=[
            {
                "type": "file_search",
                "vector_store_ids": [vector_store_id],
            }
        ],
    )

    try:
        response_dict = response.model_dump()
    except AttributeError:
        # 旧バージョンSDK互換: model_dump が無ければ辞書化 fallback
        response_dict = getattr(response, "to_dict", lambda: {})()
    answer_text = response_dict.get("output_text") or ""
    citations = extract_citations(response_dict)
    return answer_text.strip(), citations


def format_answer(answer: str, citations: List[Dict[str, Any]]) -> str:
    """回答文と引用リストを UI 向けのテキストに整形する。"""
    if not answer:
        return "該当データがありません！"

    lines = [answer]
    if citations:
        lines.append("")
        lines.append("引用:")
        for citation in citations:
            file_name = citation.get("file_name") or citation.get("file_id") or "unknown"
            quote = citation.get("quote")
            if quote:
                lines.append(f"- {quote} [{file_name}]")
            else:
                lines.append(f"- [{file_name}]")
    return "\n".join(lines)


def produce_answer(question: str, top_k: int) -> str:
    """Gradio UI からの質問を処理し、Responses API の結果を返す。"""
    question = (question or "").strip()
    if not question:
        return "質問が入力されていません。"

    try:
        vector_store_id = get_vector_store_id()
        answer, citations = answer_with_file_search(vector_store_id, question, top_k=top_k)
        return format_answer(answer, citations)
    except RuntimeError as runtime_error:
        return str(runtime_error)
    except Exception:
        return "システムエラーが発生しました。時間をおいて再度お試しください。"


def respond(question: str, top_k: int = 5) -> str:
    """Gradio の UI から呼び出されるハンドラ。"""
    try:
        return produce_answer(question, int(top_k))
    except Exception:
        return "システムエラーが発生しました。時間をおいて再度お試しください。"


with gr.Blocks() as demo:
    gr.Markdown(
        """
        ## MCP Storage レスポンダ
        OpenAI Responses API の file_search ツールを利用し、指定した Vector Store から根拠付きで回答します。
        OPENAI_API_KEY と VECTOR_STORE_ID を環境変数に設定してからご利用ください。
        """
    )

    with gr.Row():
        question_box = gr.Textbox(
            label="質問",
            placeholder="ここに質問を入力してください（例: 製品Aの対応OSは？）",
            lines=4,
        )
    with gr.Row():
        top_k_slider = gr.Slider(
            minimum=1,
            maximum=10,
            value=5,
            step=1,
            label="検索件数 (top_k)",
        )
    with gr.Row():
        output_box = gr.Textbox(label="回答", lines=10)
    with gr.Row():
        submit_button = gr.Button("検索")

    submit_button.click(respond, inputs=[question_box, top_k_slider], outputs=output_box)


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