import os
import json
from datetime import datetime

import pandas as pd
import requests
import streamlit as st

# ================== CONFIG ==================
API_URL = os.getenv("API_URL", "http://127.0.0.1:8001")

st.set_page_config(
    page_title="LearnLanguage 2026 • Tutor",
    page_icon="🧠",
    layout="wide",
)

# ================== HELPERS ==================
def api_alive() -> bool:
    """Check if backend is reachable."""
    try:
        r = requests.get(f"{API_URL}/", timeout=3)
        return r.status_code == 200
    except Exception:
        return False


def safe_post(url: str, payload: dict, headers: dict | None = None, timeout: int = 30):
    """POST helper with nice error handling."""
    try:
        r = requests.post(url, json=payload, headers=headers, timeout=timeout)
        return r
    except requests.exceptions.RequestException as e:
        st.error(f"Connection error: {e}")
        return None


def safe_get(url: str, headers: dict | None = None, timeout: int = 10):
    """GET helper with nice error handling."""
    try:
        r = requests.get(url, headers=headers, timeout=timeout)
        return r
    except requests.exceptions.RequestException as e:
        st.error(f"Connection error: {e}")
        return None


def stream_chat(message: str):
    """Call backend /chat endpoint (non-streaming in your backend, returns final JSON)."""
    payload = {"message": message, "mode": st.session_state.mode}
    headers = {"Authorization": f"Bearer {st.session_state.token}"}

    r = safe_post(f"{API_URL}/chat", payload, headers=headers, timeout=180)
    if r is None:
        raise RuntimeError("Backend not reachable.")
    r.raise_for_status()
    data = r.json()
    yield ("final", data)


# ================== SESSION STATE ==================
if "messages" not in st.session_state:
    st.session_state.messages = []

if "last" not in st.session_state:
    st.session_state.last = None

if "mode" not in st.session_state:
    st.session_state.mode = "conversation"

if "token" not in st.session_state:
    st.session_state.token = None

# ================== SIDEBAR ==================
with st.sidebar:
    st.markdown("## 🔐 Account")

    # Show backend status
    if api_alive():
        st.success("Backend: online ✅")
    else:
        st.error("Backend: offline ❌")

    tabL, tabR = st.tabs(["Login", "Register"])

    with tabL:
        email = st.text_input("Email", key="login_email")
        pwd = st.text_input("Password", type="password", key="login_pwd")

        if st.button("Login", use_container_width=True):
            r = safe_post(
                f"{API_URL}/auth/login",
                {"email": email, "password": pwd},
                timeout=20
            )

            if r is None:
                st.stop()

            if r.status_code == 200:
                try:
                    st.session_state.token = r.json()["token"]
                except Exception:
                    st.error("Login response is not valid JSON.")
                    st.stop()

                st.success("Logged in ✅")
                st.rerun()
            else:
                st.error(r.text)

    with tabR:
        email2 = st.text_input("Email", key="reg_email")
        username2 = st.text_input("Username", key="reg_username")
        pwd2 = st.text_input("Password", type="password", key="reg_pwd")

        if st.button("Create account", use_container_width=True):
            r = safe_post(
                f"{API_URL}/auth/register",
                {"email": email2, "username": username2, "password": pwd2},
                timeout=20
            )

            if r is None:
                st.stop()

            if r.status_code == 200:
                try:
                    st.session_state.token = r.json()["token"]
                except Exception:
                    st.error("Register response is not valid JSON.")
                    st.stop()

                st.success("Account created ✅")
                st.rerun()
            else:
                st.error(r.text)

    # If user logged in: show profile + logout
    if st.session_state.token:
        me = safe_get(
            f"{API_URL}/auth/me",
            headers={"Authorization": f"Bearer {st.session_state.token}"},
            timeout=10
        )

        if me is not None and me.status_code == 200:
            try:
                st.success(f"Connected as: {me.json().get('username', '—')}")
            except Exception:
                st.info("Connected, but /auth/me returned invalid JSON.")

        if st.button("Logout", use_container_width=True):
            st.session_state.token = None
            st.session_state.messages = []
            st.session_state.last = None
            st.rerun()

# Stop early if backend is down
if not api_alive():
    st.error("API backend not running (FastAPI). Check Space Logs.")
    st.stop()

# ================== MAIN ==================
st.title("🧠 LearnLanguage • Streaming Tutor")
st.caption("Streaming replies • Corrections • Exercises • Progress-ready")

if not st.session_state.token:
    st.warning("Please login first to start chatting.")
    st.stop()

# Optional: mode selector
st.session_state.mode = st.selectbox(
    "Mode",
    options=["conversation", "correction", "exercise"],
    index=["conversation", "correction", "exercise"].index(st.session_state.mode)
)

colA, colB = st.columns([5, 1])

with colA:
    user_msg = st.text_input(
        "Type your message (English)",
        placeholder="e.g., I like my city because it is calm..."
    )

with colB:
    send = st.button("Send 🚀", use_container_width=True)

# ================== SEND LOGIC ==================
if send and user_msg.strip():
    ts = datetime.now().strftime("%H:%M")

    st.session_state.messages.append({"role": "user", "text": user_msg, "ts": ts})

    placeholder = st.empty()

    try:
        for kind, data in stream_chat(user_msg):
            if kind == "text":
                # (your backend currently returns final JSON only)
                placeholder.markdown(f"**Tutor (streaming…)**\n\n{data}")
            else:
                res = data
                st.session_state.last = res
                st.session_state.messages.append({
                    "role": "bot",
                    "text": res.get("reply", ""),
                    "ts": datetime.now().strftime("%H:%M")
                })
                break

        st.rerun()

    except requests.HTTPError as e:
        st.error(f"API returned an error: {e}")
    except Exception as e:
        st.error(f"Streaming error: {e}")

# ================== CHAT HISTORY ==================
st.markdown("## 💬 Conversation")

for m in st.session_state.messages[-30:]:
    who = "You" if m["role"] == "user" else "Tutor"
    st.markdown(f"**{who} • {m['ts']}**")
    st.write(m["text"])

# ================== TUTOR PANEL ==================
st.markdown("## 🧾 Tutor Panel")

res = st.session_state.last or {}
tabs = st.tabs(["Feedback", "Exercises", "Raw JSON"])

with tabs[0]:
    st.markdown("**Corrected text**")
    st.write(res.get("corrected_text", "—"))

    st.markdown("**Corrections**")
    corrections = res.get("corrections", [])
    if corrections:
        st.dataframe(pd.DataFrame(corrections))
    else:
        st.success("No corrections 🎉")

    st.markdown("**Follow-up question**")
    st.write(res.get("followup_question", "—"))

with tabs[1]:
    st.json(res.get("exercise", {}))
    st.json(res.get("exercise_from_progress", {}))

with tabs[2]:
    st.json(res)