import streamlit as st
from groq import Groq
import requests

# =======================
# API Clients Initialization
# =======================
groq_client = None
groq_status = "❌"
if "GROQ_API_KEY" in st.secrets:
    try:
        groq_client = Groq(api_key=st.secrets["GROQ_API_KEY"])
        groq_client.models.list()  # test
        groq_status = "✅"
    except Exception:
        groq_client = None
        groq_status = "❌"

huggingface_api_key = st.secrets.get("HF_API_KEY", None)
hf_status = "❌"
if huggingface_api_key:
    try:
        headers = {"Authorization": f"Bearer {huggingface_api_key}"}
        response = requests.get("https://api-inference.huggingface.co/status", headers=headers)
        if response.status_code in [200, 401, 403]:
            hf_status = "✅"
    except Exception:
        hf_status = "❌"


# =======================
# Session State
# =======================
if "projects" not in st.session_state:
    st.session_state.projects = {"Default Chat": {"tab": "generate", "history": []}}
if "active_project" not in st.session_state:
    st.session_state.active_project = "Default Chat"


# =======================
# Sidebar Project Manager
# =======================
st.sidebar.title("📂 Chats / Projects")

# Sidebar only scrolls when hovered
st.markdown(
    """
    <style>
    [data-testid="stSidebar"] {
        overflow-y: hidden !important;   /* No scroll by default */
        height: 100vh !important;
    }
    [data-testid="stSidebar"]:hover {
        overflow-y: auto !important;    /* Scroll only when hovered */
    }
    </style>
    """,
    unsafe_allow_html=True
)

with st.sidebar:
    for project_name in list(st.session_state.projects.keys()):
        cols = st.columns([4, 1])
        with cols[0]:
            if st.button(project_name, use_container_width=True):
                st.session_state.active_project = project_name
        with cols[1]:
            with st.popover("⋮"):
                new_name = st.text_input("Rename", value=project_name, key=f"rename_{project_name}")
                if st.button("Save", key=f"save_{project_name}"):
                    if new_name and new_name != project_name:
                        st.session_state.projects[new_name] = st.session_state.projects.pop(project_name)
                        if st.session_state.active_project == project_name:
                            st.session_state.active_project = new_name
                        st.rerun()
                if st.button("Delete", key=f"delete_{project_name}"):
                    if project_name in st.session_state.projects and len(st.session_state.projects) > 1:
                        st.session_state.projects.pop(project_name)
                        st.session_state.active_project = list(st.session_state.projects.keys())[0]
                        st.rerun()

    if st.button("➕ New Chat"):
        new_name = f"Chat {len(st.session_state.projects) + 1}"
        st.session_state.projects[new_name] = {"tab": "generate", "history": []}
        st.session_state.active_project = new_name
        st.rerun()

# =======================
# Sidebar API Settings
# =======================
st.sidebar.markdown("---")
st.sidebar.title("⚙️ Settings")

st.sidebar.markdown(f"**Groq API Status:** {groq_status}")
st.sidebar.markdown(f"**HuggingFace API Status:** {hf_status}")

api_priority = st.sidebar.radio(
    "Choose API Priority",
    ["Groq First", "HuggingFace First"],
    index=0
)

model_choice = st.sidebar.selectbox(
    "Choose Model",
    ["llama-3.1-8b-instant", "llama-3.1-70b-versatile", "mixtral-8x7b-32768"]
)


# =======================
# Helper Functions
# =======================
def call_groq(system_prompt, chat_history):
    if not groq_client:
        return None
    try:
        response = groq_client.chat.completions.create(
            model=model_choice,
            messages=[{"role": "system", "content": system_prompt}]
                     + [{"role": role, "content": msg} for role, msg in chat_history],
            temperature=0.4
        )
        return response.choices[0].message.content
    except Exception as e:
        st.warning(f"⚠️ Groq API error: {e}")
        return None


def call_huggingface(system_prompt, chat_history):
    if not huggingface_api_key:
        return None
    try:
        headers = {"Authorization": f"Bearer {huggingface_api_key}"}
        payload = {
            "inputs": "\n".join([msg for _, msg in chat_history]),
            "parameters": {"temperature": 0.5, "max_new_tokens": 500}
        }
        response = requests.post(
            f"https://api-inference.huggingface.co/models/{model_choice}",
            headers=headers,
            json=payload
        )
        if response.status_code == 200:
            data = response.json()
            if isinstance(data, list) and "generated_text" in data[0]:
                return data[0]["generated_text"]
            return str(data)
        else:
            st.warning(f"⚠️ Hugging Face error: {response.text}")
            return None
    except Exception as e:
        st.warning(f"⚠️ Hugging Face exception: {e}")
        return None


def get_ai_response(system_prompt, chat_history):
    if api_priority == "Groq First":
        ai_msg = call_groq(system_prompt, chat_history)
        if ai_msg is None:
            ai_msg = call_huggingface(system_prompt, chat_history)
    else:
        ai_msg = call_huggingface(system_prompt, chat_history)
        if ai_msg is None:
            ai_msg = call_groq(system_prompt, chat_history)
    return ai_msg or "❌ Both APIs failed. Please check your API keys or try again later."


# =======================
# Main Title
# =======================
st.title("🤖 CodeCraft AI - Mini Copilot (Chat Edition)")
st.write("")  # Empty line after title


# =======================
# Chat UI
# =======================
def chat_ui(tab_name, system_prompt, input_key):
    project = st.session_state.projects[st.session_state.active_project]
    chat_history = project["history"]

    st.subheader(f"{tab_name} ({st.session_state.active_project})")

    chat_container = st.container()
    with chat_container:
        for role, msg in chat_history:
            with st.chat_message(role):
                if role == "assistant" and "```" in msg:
                    parts = msg.split("```")
                    for i, part in enumerate(parts):
                        if i % 2 == 1:
                            lang, *code_lines = part.split("\n")
                            code = "\n".join(code_lines)
                            st.code(code, language=lang if lang else "python")
                        else:
                            st.write(part)
                else:
                    st.write(msg)

    user_input = st.chat_input("Type your message...", key=input_key)
    if user_input:
        chat_history.append(("user", user_input))
        with st.spinner("🤔 Thinking..."):
            ai_msg = get_ai_response(system_prompt, chat_history)
        chat_history.append(("assistant", ai_msg))
        st.rerun()


# =======================
# Tabs
# =======================
tab1, tab2, tab3 = st.tabs(["💡 Generate Code", "🛠 Debug Code", "📘 Explain Code"])
with tab1:
    chat_ui("💡 Generate Code",
            "You are a helpful coding assistant. Generate correct code first, then a short simple explanation.",
            "generate_input")
with tab2:
    chat_ui("🛠 Debug Code",
            "You are an expert code debugger. Fix errors and give corrected code, then explain what changed and why.",
            "debug_input")
with tab3:
    chat_ui("📘 Explain Code",
            "You are a teacher that explains code in simple words. The user pastes code, and you explain step by step.",
            "explain_input")

# =======================
# Footer
# =======================
st.markdown("---")
st.caption("✨ CodeCraft may make mistakes. Always check important info.")