import streamlit as st
import subprocess
import sys
import shlex
from pathlib import Path
from PIL import Image

APP_DIR = Path(__file__).resolve().parent
PIPELINE_SCRIPT = APP_DIR / "prompt_to_card_pipeline.py"
TEXT_CLEANER_PATH = APP_DIR / "text-cleaner" / "text_cleaning_pipeline.py"
INFER_SCRIPT_PATH = APP_DIR / "clean-text-to-keywords" / "infer_json_usage.py"
CHECKPOINT_PATH = APP_DIR / "pokemon_card_lora" / "training_history.pt"
TEMPLATE_PATH = APP_DIR / "clean-text-to-keywords" / "json_template_example.json"
IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".webp", ".bmp")


def _extract_image_from_stdout(stdout: str) -> Path | None:
    for line in reversed(stdout.splitlines()):
        # Try the whole line, then the part after the last colon
        # (handles "Card generated and saved to: generated_card.png")
        raw = line.strip().strip("\"'")
        candidates = [raw]
        if ":" in raw:
            candidates.append(raw.rsplit(":", 1)[1].strip().strip("\"'"))

        for text in candidates:
            if not text:
                continue
            candidate = Path(text)
            if not candidate.is_absolute():
                candidate = APP_DIR / candidate
            if candidate.suffix.lower() in IMAGE_EXTENSIONS and candidate.exists():
                return candidate

    return None



def run_prompt_pipeline(prompt_text: str) -> tuple[Path | None, str, list[str]]:
    cmd = [
        sys.executable, "prompt_to_card_pipeline.py",
        prompt_text,
        "--text-cleaner-path", "text-cleaner/text_cleaning_pipeline.py",
        "--infer-script-path", "clean-text-to-keywords/infer_json_usage.py",
        "--checkpoint", "pokemon_card_lora",
        "--template", "clean-text-to-keywords/json_template_example.json",
        "--generator-module", "card_generator_adapter.py",
        "--device", "cuda",
        "--save-path", "generated_card.png",
        "--print-json",
    ]

    result = subprocess.run(
        cmd,
        cwd=APP_DIR,
        capture_output=True,
        text=True,
        check=False,
    )

    full_output = (result.stdout or "") + ("\n" + result.stderr if result.stderr else "")

    if result.returncode != 0:
        return None, full_output.strip() or "Erreur inconnue pendant le pipeline.", cmd

    image_path = _extract_image_from_stdout(result.stdout or "")
    return image_path, full_output.strip(), cmd

# ------------------------------------------------------------------ #
#  Configuration                                                      #
# ------------------------------------------------------------------ #

st.set_page_config(
    page_title="Générateur de Carte Pokémon",
    page_icon=Image.open(Path(__file__).with_name("pokeball.png")),
    layout="centered",
)

logo_col, title_col = st.columns([1, 6], vertical_alignment="center")

with logo_col:
    st.image(Image.open(Path(__file__).with_name("pokeball.png")), width=72)

with title_col:
    st.title("Générateur de Carte Pokémon")
st.markdown("Décrivez votre Pokémon en langage naturel et laissez la magie opérer !")

# ------------------------------------------------------------------ #
#  Saisie utilisateur                                                 #
# ------------------------------------------------------------------ #

raw_text = st.text_area(
    label="Description de votre Pokémon",
    placeholder=(
        "Ex: My pokemon is called Pyrokar! Its a huge fire dragon with massive "
        "red wings and shoots flames from its mouth... super fast n aggressive >:("
    ),
    height=180,
)

st.markdown(
    """
    <style>
    div.stButton > button {
        background-color: #d62828;
        color: white;
        border: 1px solid #b91c1c;
    }
    div.stButton > button:hover {
        background-color: #b91c1c;
        border-color: #991b1b;
        color: white;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

generate = st.button(" Générer la carte", use_container_width=True)

# ------------------------------------------------------------------ #
#  Pipeline                                                           #
# ------------------------------------------------------------------ #

if generate:
   if not raw_text.strip():
        st.warning("Veuillez entrer une description avant de générer.")
   else:
        with st.spinner("Génération de la carte Pokémon..."):
            image, logs, _cmd = run_prompt_pipeline(raw_text)

        if image is not None:
            st.image(image, caption="Carte Pokémon générée", width="stretch")
            if logs:
                with st.expander("Logs pipeline"):
                    st.code(logs)
        else:
            st.error("Aucune image générée détectée. Vérifiez les chemins du pipeline.")
            if logs:
                with st.expander("Logs pipeline"):
                    st.code(logs)