"""Streamlit chat UI - alternative front-end (Streamlit Community Cloud).

Run:  streamlit run app/streamlit_app.py
"""
from __future__ import annotations

import sys
from pathlib import Path

import streamlit as st

sys.path.append(str(Path(__file__).resolve().parents[1]))


@st.cache_resource
def get_assistant():
    from src.rag.generator import CodeAssistant
    return CodeAssistant.from_config(with_index=True)


st.set_page_config(page_title="Code Generation Assistant", page_icon="</>")
st.title("Code Generation Assistant")

with st.sidebar:
    mode = st.radio("Mode", ["rag", "baseline"], index=0)
    use_repair = st.checkbox("Agentic repair (run + self-fix)", value=False)

assistant = get_assistant()

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

for m in st.session_state.messages:
    with st.chat_message(m["role"]):
        st.markdown(m["content"])

if prompt := st.chat_input("Describe the function you want..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"), st.spinner("Generating..."):
        if use_repair:
            from src.agent.repair_loop import make_repair_generator, repair_loop
            gen_fn = make_repair_generator(assistant)
            trace = repair_loop(prompt, gen_fn, check_program_fn=lambda c: c, max_iters=3)
            code = trace.final_code
            st.code(code, language="python")
            st.caption(f"Repair: {'passed' if trace.success else 'best effort'} in {trace.iterations} iters")
            answer = code
        else:
            result = assistant.generate(prompt, mode=mode, return_sources=True)
            if isinstance(result, tuple):
                code, sources = result
                st.code(code, language="python")
                with st.expander("Retrieved examples"):
                    for s in sources:
                        st.write(f"({s['score']:.2f}) {s['docstring']}")
            else:
                code = result
                st.code(code, language="python")
            answer = code
    st.session_state.messages.append({"role": "assistant", "content": f"```python\n{answer}\n```"})