"""Gradio API adapter for the ASU CodonTransformer wiki."""

from __future__ import annotations

import importlib
import os
from pathlib import Path

import gradio as gr
import spaces
from huggingface_hub import hf_hub_download

from api.huggingface import register_adp1_organism
from api.service import (
    MAX_FASTA_BYTES,
    CodonTransformerService,
    _ensure_codontransformer_import,
    parse_dna_text,
)

_ensure_codontransformer_import()
codon_utils = importlib.import_module("CodonTransformer.CodonUtils")
register_adp1_organism(codon_utils)

model_path = os.environ.get("CODONTRANSFORMER_MODEL_PATH")
model_repo_id = os.environ.get("CODONTRANSFORMER_MODEL_REPO_ID", "").strip()
if not model_path and model_repo_id:
    model_path = hf_hub_download(
        repo_id=model_repo_id,
        filename=os.environ.get(
            "CODONTRANSFORMER_MODEL_FILENAME",
            "a_baylyi_adp1/a_baylyi_adp1.ckpt",
        ),
    )
    os.environ["CODONTRANSFORMER_MODEL_PATH"] = model_path
elif not model_path:
    bundled_checkpoint = (
        Path(__file__).parent
        / "checkpoints"
        / "a_baylyi_adp1"
        / "a_baylyi_adp1.ckpt"
    )
    if bundled_checkpoint.is_file():
        model_path = str(bundled_checkpoint)
        os.environ["CODONTRANSFORMER_MODEL_PATH"] = model_path

predictor = CodonTransformerService()


def list_organisms() -> list[dict[str, int | str]]:
    return predictor.supported_organisms


@spaces.GPU(duration=90)
def predict(organism: str, sequence_text: str) -> dict[str, list[dict[str, str]]]:
    if len(sequence_text.encode("utf-8")) > MAX_FASTA_BYTES:
        raise ValueError("The FASTA input must be 25 MB or smaller.")
    records = parse_dna_text(sequence_text)
    return {"results": predictor.predict(records, organism)}


with gr.Blocks(title="ASU CodonTransformer") as demo:
    gr.Markdown(
        "# ASU CodonTransformer\n"
        "Optimize coding DNA for an organism supported by the loaded model."
    )
    organism_input = gr.Textbox(label="Target organism")
    sequence_input = gr.Textbox(
        label="DNA sequence or FASTA",
        lines=8,
        placeholder=">gene1\nATGAAACCCGGGTAA",
    )
    organisms_output = gr.JSON(visible=False)
    gr.Button("List supported organisms").click(
        fn=list_organisms,
        inputs=[],
        outputs=organisms_output,
        api_name="organisms",
    )
    prediction_output = gr.JSON(label="Prediction")
    gr.Button("Optimize DNA").click(
        fn=predict,
        inputs=[organism_input, sequence_input],
        outputs=prediction_output,
        api_name="predict",
    )

demo.queue()
demo.launch(
    server_name="0.0.0.0",
    server_port=int(os.environ.get("PORT", "7860")),
    show_error=True,
)