"""Training page: start model training and poll for status."""

import time

import streamlit as st

from frontend import api_client


def show_training_page():
    project = st.session_state.get("current_project")
    if not project:
        st.warning("No project selected")
        return

    st.title(f"Train Model: {project['name']}")

    if st.button("< Back to Dashboard"):
        st.session_state["page"] = "dashboard"
        st.rerun()

    # Check for existing runs
    resp = api_client.get(f"/projects/{project['id']}/runs/latest")
    has_existing = resp.status_code == 200
    existing_run = resp.json() if has_existing else None

    if has_existing and existing_run.get("status") in ("pending", "running"):
        _show_progress(project["id"], existing_run)
        return

    if has_existing and existing_run.get("status") == "completed":
        st.success(
            f"Last training completed in {existing_run.get('duration_secs', 0):.0f}s"
        )
        if st.button("View Results"):
            st.session_state["page"] = "results"
            st.rerun()
        st.divider()

    if has_existing and existing_run.get("status") == "failed":
        st.error(f"Last training failed: {existing_run.get('error_message', 'Unknown error')}")

    # Start training
    st.subheader("Start New Training Run")
    st.markdown(
        "This will fit a Bayesian Marketing Mix Model on your data. "
        "Typical training time: **5-15 minutes** depending on data size."
    )

    if st.button("Train Model", type="primary"):
        resp = api_client.post(f"/projects/{project['id']}/train")
        if resp.status_code == 202:
            st.session_state["current_run_id"] = resp.json()["run_id"]
            st.rerun()
        else:
            st.error(resp.json().get("detail", "Failed to start training"))


def _show_progress(project_id: str, run: dict):
    """Poll for training progress."""
    status = run.get("status", "pending")
    stage_labels = {
        "pending": "Queued — waiting to start...",
        "running": "Training model — this may take 5-15 minutes...",
    }

    st.info(stage_labels.get(status, "Processing..."))
    progress_bar = st.progress(0.3 if status == "pending" else 0.6)

    placeholder = st.empty()
    placeholder.text("Checking status...")

    # Auto-refresh
    time.sleep(5)
    resp = api_client.get(f"/projects/{project_id}/runs/latest")
    if resp.status_code == 200:
        updated = resp.json()
        if updated.get("status") == "completed":
            progress_bar.progress(1.0)
            st.success("Training complete!")
            if st.button("View Results"):
                st.session_state["page"] = "results"
                st.rerun()
        elif updated.get("status") == "failed":
            st.error(f"Training failed: {updated.get('error_message')}")
        else:
            st.rerun()
