# Flyto2 Core - Source Available License
# Copyright (c) 2025 Flyto2. All Rights Reserved.
#
# This source code is licensed under the Flyto2 Source Available License v1.0.
# Commercial use requires a license. See LICENSE for details.

"""
HuggingFace Translation Module

Translate text between languages.
"""
import logging
from typing import Any, Dict

from ...registry import register_module
from ...schema import compose, presets
from .constants import TaskType, ModuleDefaults, ModuleColors, Subcategory
from ._base import HuggingFaceTaskExecutor, normalize_text_result


logger = logging.getLogger(__name__)

_executor = HuggingFaceTaskExecutor(TaskType.TRANSLATION)


@register_module(
    module_id='huggingface.translation',
    stability="beta",
    version=ModuleDefaults.VERSION,
    category=ModuleDefaults.CATEGORY,
    subcategory=Subcategory.TEXT,
    tags=['huggingface', 'text', 'translation', 'nlp', 'language'],
    label='Translation',
    label_key='modules.huggingface.translation.label',
    description='Translate text between languages using Helsinki-NLP, mBART, etc.',
    description_key='modules.huggingface.translation.description',
    icon='Languages',
    color=ModuleColors.TRANSLATION,

    input_types=['text'],
    output_types=['text'],
    can_connect_to=['string.*'],
    can_receive_from=['data.*', 'string.*', 'file.*', 'http.*', 'flow.*'],

    timeout=ModuleDefaults.TIMEOUT,
    retryable=ModuleDefaults.RETRYABLE,
    max_retries=ModuleDefaults.MAX_RETRIES,
    concurrent_safe=ModuleDefaults.CONCURRENT_SAFE,

    requires_credentials=ModuleDefaults.REQUIRES_CREDENTIALS,
    handles_sensitive_data=ModuleDefaults.HANDLES_SENSITIVE_DATA,

    # Schema-driven params
    params_schema=compose(
        presets.HF_MODEL_ID(task=TaskType.TRANSLATION),
        presets.HF_TEXT_INPUT(),
        presets.HF_SOURCE_LANG(),
        presets.HF_TARGET_LANG(),
    ),
    output_schema={
        'translation_text': {'type': 'string', 'description': 'Translated text',
                'description_key': 'modules.huggingface.translation.output.translation_text.description'}
    },
    author=ModuleDefaults.AUTHOR,
    license=ModuleDefaults.LICENSE,
    timeout_ms=30000,
)
async def huggingface_translation(context: Dict[str, Any]) -> Dict[str, Any]:
    """Translate text using HuggingFace models"""
    params = context['params']
    model_id = params['model_id']
    text = params['text']

    kwargs = {}
    if params.get('source_lang'):
        kwargs['src_lang'] = params['source_lang']
    if params.get('target_lang'):
        kwargs['tgt_lang'] = params['target_lang']

    exec_result = await _executor.execute(
        model_id=model_id,
        inputs=text,
        **kwargs
    )

    translation_text = normalize_text_result(exec_result['raw_result'])
    logger.info(f"Translated to {len(translation_text)} characters")

    return {
        'ok': True,
        'translation_text': translation_text,
        'model_id': model_id,
        'runtime': exec_result['runtime']
    }
