import streamlit as st
from egritools import lmstudio
from datetime import datetime
import json
from pathlib import Path

LOG_PATH = Path("chat_log.json")

def save_chat_log(entry):
    try:
        if LOG_PATH.exists():
            existing = json.loads(LOG_PATH.read_text(encoding="utf-8"))
        else:
            existing = []
        existing.append(entry)
        LOG_PATH.write_text(json.dumps(existing, indent=2, ensure_ascii=False), encoding="utf-8")
    except Exception as e:
        st.error(f"Failed to save chat log: {e}")

def code_chat_interface():
    st.title("💻 Code Chat Assistant")

    try:
        models = lmstudio.get_models()
    except Exception as e:
        st.error(f"Could not fetch models: {e}")
        return

    if not models:
        st.warning("No models available from LM Studio.")
        return

    model = st.selectbox("Choose a model", models)

    # Optional model settings
    col1, col2 = st.columns(2)
    with col1:
        temperature = st.slider("Temperature", 0.0, 1.5, 0.1, step=0.1)
    with col2:
        top_p = st.slider("Top-p", 0.1, 1.0, 0.9, step=0.1)

    max_tokens = st.slider("Max Tokens", 128, 4096, 1024, step=128)

    use_streaming = st.checkbox("Use streaming mode", value=False)

    st.markdown("### ✏️ Code Context (optional)")
    code_input = st.text_area("Paste code you want the model to consider", height=200, placeholder="Paste your code here...")

    st.markdown("### ❓ Your Question")
    question = st.text_area("Ask a question about the code or programming in general", height=150, placeholder="e.g. What does this function do?")

    if st.button("🧠 Ask"):
        if not question.strip():
            st.warning("Please enter a question.")
            return

        # Build the final prompt
        full_prompt = f"""You are an expert programmer. Analyze the code and answer the user's question.

Code:
{code_input.strip() if code_input else '[no code provided]'}

Question:
{question.strip()}
"""

        with st.spinner("Generating answer..."):
            if use_streaming:
                st.markdown("**Response:**")
                response_container = st.empty()
                collected = ""
                for chunk in lmstudio.stream_chat_with_model(model, full_prompt):
                    collected += chunk
                    response_container.markdown("```python\n" + collected.strip() + "\n```")
                answer = collected.strip()
            else:
                answer = lmstudio.chat_with_model(model, full_prompt)
                st.markdown("**Response:**")
                st.code(answer.strip(), language="python")

        # Log interaction
        log_entry = {
            "timestamp": datetime.utcnow().isoformat(),
            "model": model,
            "temperature": temperature,
            "top_p": top_p,
            "max_tokens": max_tokens,
            "code": code_input.strip(),
            "question": question.strip(),
            "prompt": full_prompt.strip(),
            "response": answer.strip()
        }
        save_chat_log(log_entry)