"""
文本 Gradio 部署加载辅助函数。

这里绕开训练 Pipeline 的数据源阶段，直接从仓库内的部署模型和词表创建文本推理资源。
三个文本 Gradio 页面会复用这些函数，避免 HF Space 启动时访问本地训练数据目录。
"""

import tensorflow as tf

from deep_learning.data.spec import TokenizerBundle
from deep_learning.data.poetry.tokenizer import load_vectorizer
from deep_learning.data.wiki.tokenizer import sentence_piece
from deep_learning.env.resolve import resolve_saved
from deep_learning.models.mini_gpt import GptModelBuilder
from deep_learning.models.rnn import RNNModelBuilder
from deep_learning.pipeline.services.model_loader import load_inference_artifact
from deep_learning.pipeline.impl.text_generation import (
    TextInferenceBundle,
    top_k
)


def _empty_docs_ds() -> tf.data.Dataset:
    return tf.data.Dataset.from_tensor_slices(tf.constant([], dtype=tf.string))


def _wiki_tokenizer_bundle() -> TokenizerBundle:
    vocab_path = resolve_saved("vocab/sentencepiece/vocabulary.proto")
    tokenizer, end_of_text, decode = sentence_piece()
    return TokenizerBundle(
        tokenizer=tokenizer,
        decode=decode,
        end_of_text=end_of_text,
        vocab_size=tokenizer.vocabulary_size(),
        vocab_path=str(vocab_path)
    )


def _poetry_tokenizer_bundle(sequence_length: int) -> TokenizerBundle:
    vocab_path = resolve_saved("vocab/poetry/vocab.txt")
    tokenizer = load_vectorizer(vocab_path, sequence_length + 1)
    vocab = tokenizer.get_vocabulary()
    end_of_text = vocab.index("$")

    def decode(token_ids: list[int]) -> str:
        chars = [
            vocab[token_id] for token_id in token_ids if token_id < len(vocab)
        ]
        return "".join(chars)

    return TokenizerBundle(
        tokenizer=tokenizer,
        decode=decode,
        end_of_text=end_of_text,
        vocab_size=len(vocab),
        vocab_path=str(vocab_path)
    )


def _text_inference_bundle(tokenizer_bundle: TokenizerBundle, max_length: int) -> TextInferenceBundle:
    return TextInferenceBundle(
        tokenizer_bundle=tokenizer_bundle,
        docs_ds=_empty_docs_ds(),
        max_length=max_length,
        sample_fn=top_k
    )


def load_wiki_gpt_deployment_inference_resource() -> TextInferenceBundle:
    return _text_inference_bundle(_wiki_tokenizer_bundle(), 256)


def load_poetry_gpt_deployment_inference_resource() -> TextInferenceBundle:
    return _text_inference_bundle(_poetry_tokenizer_bundle(100), 100)


def load_poetry_rnn_deployment_inference_resource() -> TextInferenceBundle:
    return _text_inference_bundle(_poetry_tokenizer_bundle(100), 100)


def load_wiki_gpt_deployment_inference_artifact(checkpoint_rule: dict) -> tuple:
    model_builder = GptModelBuilder(
        hidden_dim=512,
        intermediate_dim=2056,
        num_heads=8,
        num_layers=8
    )
    return load_inference_artifact(
        checkpoint_rule,
        model_builder.load_inference_artifact,
        resource_factory=load_wiki_gpt_deployment_inference_resource
    )


def load_poetry_gpt_deployment_inference_artifact(checkpoint_rule: dict) -> tuple:
    model_builder = GptModelBuilder(
        hidden_dim=512,
        intermediate_dim=2056,
        num_heads=4,
        num_layers=4
    )
    return load_inference_artifact(
        checkpoint_rule,
        model_builder.load_inference_artifact,
        resource_factory=load_poetry_gpt_deployment_inference_resource
    )


def load_poetry_rnn_deployment_inference_artifact(checkpoint_rule: dict) -> tuple:
    model_builder = RNNModelBuilder(num_layers=2, embedding_dim=100, hidden_dim=512)
    return load_inference_artifact(
        checkpoint_rule,
        model_builder.load_inference_artifact,
        resource_factory=load_poetry_rnn_deployment_inference_resource
    )