import gc
import os

import gradio as gr
import torch

from qwen_asr import Qwen3ASRModel


MODEL_ID = "Cartik/Sonexa-1-ASR"

# ============================================================
# DEVICE
# ============================================================

IS_ZERO_GPU = os.environ.get("SPACES_ZERO_GPU") == "true"
HAS_CUDA = torch.cuda.is_available()

USE_GPU = IS_ZERO_GPU or HAS_CUDA

if USE_GPU:
    DEVICE = "cuda"
    DTYPE = torch.bfloat16
else:
    DEVICE = "cpu"
    DTYPE = torch.float32

print("=" * 60)
print("Sonexa ASR")
print("=" * 60)
print(f"Model: {MODEL_ID}")
print(f"ZeroGPU: {IS_ZERO_GPU}")
print(f"CUDA: {HAS_CUDA}")
print(f"Device: {DEVICE}")
print(f"Dtype: {DTYPE}")
print("=" * 60)


# ============================================================
# GRADIO CLIENT PATCH
# ============================================================
#
# Некоторые версии gradio_client падают на JSON Schema,
# где additionalProperties / другие поля имеют значение bool.
#
# Без этого show_api может падать с:
#
# TypeError: argument of type 'bool' is not iterable
#
# В результате API predict становится недоступным.
# ============================================================

try:
    import gradio_client.utils as gc_utils

    original_json_schema = gc_utils._json_schema_to_python_type

    def patched_json_schema_to_python_type(schema, defs=None):
        if isinstance(schema, bool):
            return "Any"

        return original_json_schema(schema, defs)

    gc_utils._json_schema_to_python_type = (
        patched_json_schema_to_python_type
    )

    print("Gradio JSON Schema patch applied.")

except Exception as e:
    print(f"Gradio patch warning: {e}")


# ============================================================
# GPU DECORATOR
# ============================================================

if USE_GPU:

    try:
        import spaces

        gpu_decorator = spaces.GPU(duration=60)

        print("ZeroGPU decorator enabled.")

    except ImportError:

        print(
            "Package 'spaces' not found. "
            "GPU decorator disabled."
        )

        def gpu_decorator(fn):
            return fn

else:

    def gpu_decorator(fn):
        return fn


# ============================================================
# LOAD MODEL
# ============================================================

print("Loading Sonexa ASR model...")

model = Qwen3ASRModel.from_pretrained(
    MODEL_ID,
    dtype=DTYPE,
    device_map="auto" if USE_GPU else "cpu",
)

print("Model loaded successfully!")


# ============================================================
# TRANSCRIBE
# ============================================================

@gpu_decorator
def transcribe(audio, language):

    if audio is None:
        raise gr.Error(
            "Загрузите аудио или запишите его с микрофона."
        )

    try:

        print(f"Audio: {audio}")
        print(f"Language: {language}")

        with torch.inference_mode():

            results = model.transcribe(
                audio=audio,
                language=language,
            )

        if isinstance(results, list):

            if not results:
                return ""

            result = results[0]

            text = getattr(result, "text", None)

            if text is not None:
                return str(text)

            if isinstance(result, str):
                return result

            if isinstance(result, dict):

                if "text" in result:
                    return str(result["text"])

                if "transcription" in result:
                    return str(result["transcription"])

            return str(result)

        if isinstance(results, str):
            return results

        text = getattr(results, "text", None)

        if text is not None:
            return str(text)

        if isinstance(results, dict):

            if "text" in results:
                return str(results["text"])

            if "transcription" in results:
                return str(results["transcription"])

        return str(results)

    except Exception as e:

        print("=" * 60)
        print("ASR ERROR")
        print(e)
        print("=" * 60)

        raise gr.Error(
            f"Ошибка распознавания: {e}"
        )

    finally:

        gc.collect()

        if HAS_CUDA:
            torch.cuda.empty_cache()


# ============================================================
# UI
# ============================================================

with gr.Blocks(
    title="Sonexa ASR"
) as demo:

    gr.Markdown("# Sonexa ASR")

    if USE_GPU:
        gr.Markdown(
            "⚡ Распознавание выполняется на GPU."
        )
    else:
        gr.Markdown(
            "⚠️ GPU недоступен — используется CPU."
        )

    with gr.Row():

        with gr.Column():

            audio = gr.Audio(
                label="Аудио",
                sources=[
                    "upload",
                    "microphone",
                ],
                type="filepath",
            )

            language = gr.Dropdown(
                choices=[
                    "Russian",
                    "English",
                    "Chinese",
                    "Japanese",
                    "Korean",
                    "German",
                    "French",
                    "Spanish",
                ],
                value="Russian",
                label="Язык",
            )

            button = gr.Button(
                "Распознать",
                variant="primary",
            )

        with gr.Column():

            output = gr.Textbox(
                label="Распознанный текст",
                lines=10,
            )

    button.click(
        fn=transcribe,
        inputs=[
            audio,
            language,
        ],
        outputs=output,
        api_name="predict",
    )


# ============================================================
# QUEUE
# ============================================================

demo.queue(
    max_size=32,
    default_concurrency_limit=1,
)


# ============================================================
# LAUNCH
# ============================================================

demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    show_api=True,
    ssr=False,
)