"""Chat view — renders both the normal chat interface and the Coding Agent mode."""

import asyncio
from pathlib import Path

import streamlit as st
from backend.managers.chat_manager import ChatManager
from backend.managers.system_prompter import SystemPrompter
from backend.agent.coding_agent import CodingAgent

from backend.managers.debug_logger import get_logger
logger = get_logger(__name__)

import asyncio
import json


# ── Agent Mode helpers ────────────────────────────────────────────────────────
def _run_async(coro):
    """Execute an async coroutine from synchronous Streamlit code.

    Streamlit runs in a synchronous context, but the CodingAgent uses async
    methods (for MCP tool calls).  This helper bridges the gap by reusing an
    already-running event loop when one exists, or creating a new one otherwise.

    Args:
        coro: The coroutine to run.

    Returns:
        The return value of the coroutine.
    """
    # REVIEW: asyncio.get_running_loop() always raises RuntimeError in a Streamlit context;
    # the try branch is dead code. The except branch always runs.
    try:
        # Reuse the loop that is already running (e.g. inside pytest-asyncio).
        loop = asyncio.get_running_loop()
    except RuntimeError:
        # No running loop in this thread — create a fresh one.
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    return loop.run_until_complete(coro)

def _start_agent(task: str):
    """Create a new CodingAgent, feed it the task, and propose the first action.
    Stores the agent and its state in session_state so Streamlit can reference
    them across reruns without losing progress.
    """
    logger.info("Starting coding agent.")
    agent = CodingAgent()
    agent.start_task(task)
    action = _run_async(agent.propose_next_action())
    st.session_state.coding_agent = agent
    st.session_state.agent_pending_action = action
    st.session_state.agent_status = "waiting_approval"
    st.session_state.agent_log = []


def _approve_action():
    """Execute the pending action, log it, then immediately propose the next step."""
    agent = st.session_state.coding_agent
    pending = st.session_state.agent_pending_action

    result = _run_async(agent.approve())
    logger.info("Approve action and propose next step.")

    # Append a record to the log so the user can review every completed step.
    st.session_state.agent_log.append({
        "thought":   pending.get("thought", ""),
        "tool":      result["tool"],
        "arguments": result.get("arguments", {}),
        "result":    result["result"],
    })

    if result["is_done"]:
        # Agent called the "done" tool — task is fully complete.
        st.session_state.agent_status = "done"
        st.session_state.agent_pending_action = None
    else:
        next_action = _run_async(agent.propose_next_action())
        st.session_state.agent_pending_action = next_action
        st.session_state.agent_status = "waiting_approval"


def _reject_action(feedback: str):
    """Reject the pending action with feedback so the agent replans.
    The pending action is discarded; the agent receives the user's feedback and
    proposes a different approach on the next call to propose_next_action().
    """
    logger.info("Rejecting proposed action.")
    agent = st.session_state.coding_agent
    agent.reject(feedback or "Please try a different approach.")
    next_action = _run_async(agent.propose_next_action())
    st.session_state.agent_pending_action = next_action
    st.session_state.agent_status = "waiting_approval"


def _followup_agent(question: str):
    """Continue a finished task by injecting a follow-up question and resuming the loop."""
    logger.info("Asking follow up question")
    agent = st.session_state.coding_agent
    agent.follow_up(question)
    action = _run_async(agent.propose_next_action())
    st.session_state.agent_pending_action = action
    st.session_state.agent_status = "waiting_approval"


def _reset_agent():
    """Clear all agent state and return to the idle (task input) screen."""
    logger.info("Resetting Agent")
    st.session_state.coding_agent = None
    st.session_state.agent_status = "idle"
    st.session_state.agent_log = []
    st.session_state.agent_pending_action = None


# ── Agent Mode UI ─────────────────────────────────────────────────────────────

def _render_arguments(args: dict):
    if not args:
        return

    with st.expander("📦 Arguments", expanded=False):

        if args.get("path"):
            st.markdown("##### 📁 Path")
            st.code(args["path"])

        if args.get("dir_path"):
            st.markdown("##### 🌳 Directory")
            st.code(args["dir_path"])

        if args.get("query"):
            st.markdown("##### 🔎 Query")
            st.code(args["query"])

        if args.get("url"):
            st.markdown("##### 🌐 URL")
            st.code(args["url"])

        if args.get("content"):
            st.markdown("##### 📝 Content")
            st.code(args["content"])

        if args.get("code"):
            st.markdown("##### 🐍 Python Code")
            st.code(args["code"], language="python")

        if args.get("max_results") is not None:
            st.markdown("##### 📊 Max Results")
            st.code(str(args["max_results"]))

        known_keys = {
            "path",
            "dir_path",
            "query",
            "content",
            "url",
            "code",
            "max_results",
        }

        extra_args = {
            k: v for k, v in args.items()
            if k not in known_keys
        }

        if extra_args:
            st.markdown("##### ⚙️ Other")
            st.code(
                json.dumps(extra_args, indent=2),
                language="json"
            )

def render_agent_mode():
    """Render the step-by-step agent UI.

    Three distinct screens based on agent_status:
      - "idle"             → task description input + Start button
      - "waiting_approval" → show proposed action, Approve / Reject / Abort
      - "done"             → success message, follow-up input, New Task button
    """
    logger.info("Agent mode.")
    # The toggle must always render so Streamlit keeps agent_mode=True in session_state.
    st.toggle("Agent Mode", key="agent_mode")

    agent_status = st.session_state.get("agent_status", "idle")
    agent_log    = st.session_state.get("agent_log", [])

    # ── Agent Log ────────────────────────────────────────────────────────────
    # Collapsed by default so it doesn't clutter the UI during active tasks.
    if agent_log:
        with st.expander(f"Agent Log — {len(agent_log)} step(s) completed", expanded=False):
            for i, step in enumerate(agent_log):
                with st.chat_message("assistant"):
                    st.markdown(f"**Step {i + 1} — `{step['tool']}`**")
                    st.caption(f"Thought: {step['thought']}")
                    if step.get("arguments"):                    
                        _render_arguments(step["arguments"])
                    result_text = step.get("result", "")
                    # Colour the result based on whether the tool succeeded or failed.
                    if result_text.startswith("ERROR") or result_text.startswith("SYNTAX ERROR"):
                        st.error(result_text)
                    elif result_text.startswith("OK") or result_text.startswith("DONE"):
                        st.success(result_text)
                    else:
                        st.code(result_text, language=None)

    # ── Idle: task input ──────────────────────────────────────────────────────
    if agent_status == "idle":
        task = st.text_area(
            "Describe what the agent should do:",
            key="agent_task_input",
            height=120,
            placeholder="e.g. Write a function that sorts a list and saves it to sorted.py",
        )
        if st.button("Start Agent", type="primary", use_container_width=True):
            if task.strip():
                with st.spinner("Agent is thinking..."):
                    _start_agent(task.strip())
                st.rerun()
            else:
                st.warning("Please describe a task first.")

    # ── Waiting: show proposed action + Approve / Reject ─────────────────────
    elif agent_status == "waiting_approval":
        pending = st.session_state.get("agent_pending_action", {})

        with st.status("Agent proposes the following step:", expanded=True):
            st.markdown(f"**Thought:** {pending.get('thought', '—')}")
            st.markdown(f"**Tool:** `{pending.get('tool', '—')}`")

            args = pending.get("arguments", {})
            if args:
                _render_arguments(args)

        feedback = st.text_input(
            "Rejection feedback (optional):",
            key="agent_reject_feedback",
            placeholder="e.g. Use a different approach...",
        )

        col1, col2, col3 = st.columns([3, 2, 2])
        with col1:
            if st.button("Approve", type="primary", use_container_width=True):
                with st.spinner("Executing and planning next step..."):
                    _approve_action()
                st.rerun()
        with col2:
            if st.button("Reject", use_container_width=True):
                with st.spinner("Agent is replanning..."):
                    _reject_action(feedback)
                st.rerun()
        with col3:
            if st.button("Abort Task", use_container_width=True):
                _reset_agent()
                st.rerun()

    # ── Done ─────────────────────────────────────────────────────────────────
    elif agent_status == "done":
        last_result = agent_log[-1]["result"] if agent_log else ""
        st.success(f"Task completed! {last_result}")

        st.divider()

        followup = st.text_area(
            "Follow-up question or correction:",
            key="agent_followup_input",
            height=80,
            placeholder="e.g. The output is wrong — it should sort descending. Can you fix that?",
        )

        col1, col2 = st.columns(2)
        with col1:
            if st.button("Ask Follow-up", type="primary", use_container_width=True):
                if followup.strip():
                    with st.spinner("Agent is thinking..."):
                        _followup_agent(followup.strip())
                    st.rerun()
                else:
                    st.warning("Please enter a follow-up question first.")
        with col2:
            if st.button("New Task", use_container_width=True):
                _reset_agent()
                st.rerun()


# ── Normal Chat helpers ───────────────────────────────────────────────────────

def _build_file_context() -> dict | None:
    """Return file context for the system prompt if a file is open and context is enabled.

    Reads from files_content cache first; falls back to FileManager if the file
    has not been loaded into the editor yet.
    """
    if not st.session_state.get("include_file_context", True):
        return None
    active_file = st.session_state.get("active_file")
    if not active_file:
        return None
    content = st.session_state.get("files_content", {}).get(active_file, "")
    if not content:
        try:
            from backend.managers.file_manager import FileManager
            fm = FileManager()
            content = fm.read_file(Path(active_file)) or ""
        except Exception:
            return None
    return {"name": Path(active_file).name, "content": content}


@st.dialog("Clear Chat")
def _clear_chat_dialog():
    """Confirmation dialog before wiping the full conversation history."""
    st.warning("All messages will be deleted. This cannot be undone.")
    col1, col2 = st.columns(2)
    with col1:
        if st.button("Clear", type="primary", use_container_width=True):
            st.session_state.chat_manager.clear_history()
            st.session_state.chat_history = []
            st.rerun()
    with col2:
        if st.button("Cancel", use_container_width=True):
            st.rerun()


# ── Normal Chat ───────────────────────────────────────────────────────────────

def _render_search_panel():
    """Render the collapsible web search panel above the chat history.

    Stores results in session_state.search_results so they are automatically
    injected as context into the next message the user sends.
    """
    search_results = st.session_state.get("search_results", [])
    label = f"🔍 Web Search  ({len(search_results)} result{'s' if len(search_results) != 1 else ''} active)" if search_results else "🔍 Web Search"

    with st.expander(label, expanded=False):
        col_input, col_btn = st.columns([5, 1])
        with col_input:
            query = st.text_input(
                "Search query",
                key="search_query_input",
                placeholder="e.g. Python asyncio best practices",
                label_visibility="collapsed",
            )
        with col_btn:
            search_clicked = st.button("Search", use_container_width=True)

        if search_clicked and query.strip():
            with st.spinner("Searching..."):
                sm = SearchManager()
                results = sm.perform_search(query.strip())
            if results:
                st.session_state.search_results = results
                st.rerun()
            else:
                st.warning("No results found.")

        # Display active results with a clear button.
        if search_results:
            st.caption("Results will be injected as context into your next message.")
            for r in search_results:
                st.markdown(f"**{r['title']}**  \n{r['snippet']}  \n[{r['url']}]({r['url']})")
                st.divider()
            if st.button("Clear search results", use_container_width=True):
                st.session_state.search_results = []
                st.rerun()


def render_normal_chat():
    """Render the standard multi-turn chat interface.

    On the first message the system prompt is injected into the history,
    including any active search results as context.
    Each subsequent message appends to the same conversation so the AI retains
    full context throughout the session. If search results are active when the
    user sends a message, they are prepended to that message as a context block.
    """
    logger.info("Chat mode")
    # Replay the conversation history as chat bubbles (skip system messages).
    for message in st.session_state.chat_history:
        if message["role"] == "system":
            continue
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Chat input — Enter to send, no extra button needed.
    # Supports /search <query> and /search clear as special commands.
    user_input = st.chat_input("Type a message or /search <query>...")
    if user_input:
        stripped = user_input.strip()

        # ── /search command ───────────────────────────────────────────────────
        if stripped.lower().startswith("/search"):
            arg = stripped[len("/search"):].strip()

            with st.chat_message("user"):
                st.markdown(stripped)

            if arg.lower() == "clear" or arg == "":
                # /search clear (or bare /search) — remove active results.
                st.session_state.search_results = []
                with st.chat_message("assistant"):
                    st.markdown("Search context cleared.")
                st.session_state.chat_history.append({"role": "user", "content": stripped})
                st.session_state.chat_history.append({"role": "assistant", "content": "Search context cleared."})
            else:
                # /search <query> — run search and store results in context.
                with st.chat_message("assistant"):
                    with st.spinner(f'Searching for "{arg}"...'):
                        sm = SearchManager()
                        results = sm.perform_search(arg)

                    if results:
                        st.session_state.search_results = results
                        summary = f"Found {len(results)} result(s) for **{arg}**. They are now in context for this chat session.\n\n"
                        for i, r in enumerate(results, 1):
                            summary += f"**{i}. [{r['title']}]({r['url']})**  \n{r['snippet']}\n\n"
                        st.markdown(summary)
                        response_text = summary
                    else:
                        msg = f'No results found for "{arg}".'
                        st.warning(msg)
                        response_text = msg

                st.session_state.chat_history.append({"role": "user", "content": stripped})
                st.session_state.chat_history.append({"role": "assistant", "content": response_text})

            st.rerun()
            return

        # ── Normal chat message ───────────────────────────────────────────────
        chat_manager = st.session_state.chat_manager
        search_results = st.session_state.get("search_results", [])

        # 5g — System-prompt logic: inject on first message, update on file change.
        if not chat_manager.get_history():
            system_prompt = SystemPrompter.generate_prompt()
            chat_manager.add_message("system", system_prompt)

        # If search results are active, prepend them as a context block so the
        # AI can reference them regardless of where in the conversation we are.
        if search_results:
            context_block = "<search_context>\n"
            for r in search_results:
                context_block += (
                    f"Title: {r['title']}\n"
                    f"URL: {r['url']}\n"
                    f"Snippet: {r['snippet']}\n\n"
                )
            context_block += "</search_context>\n\n"
            message_to_send = context_block + user_input
        else:
            message_to_send = user_input

        # Show the original user text in the UI (not the context-enriched version).
        with st.chat_message("user"):
            st.markdown(user_input)

        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                try:
                    ai_response = chat_manager.send_message(message_to_send)
                except Exception as e:
                    ai_response = f"Error: {e}"
            st.markdown(ai_response)

        st.session_state.chat_history.append({"role": "user",      "content": user_input})
        st.session_state.chat_history.append({"role": "assistant", "content": ai_response})
        st.rerun()

    # 5d — Clear Chat opens a confirmation dialog instead of deleting immediately.
    if st.button("🗑️ Clear Chat"):
        _clear_chat_dialog()

    # REVIEW: duplicate widget key — "agent_mode" toggle is already rendered inside render_agent_mode();
    # having two st.toggle calls with the same key on the same page will raise a DuplicateWidgetID error.
    st.toggle("Agent Mode", key="agent_mode")

    # 5h — Settings expander: file context toggle, model, token limit, custom prompt.
    with st.expander("⚙️ Settings", expanded=False):
        st.toggle("Include current file as context", key="include_file_context", value=True)

        st.divider()

        default_model = chat_manager.model or ""
        model_options = [default_model] if default_model else []
        for m in ["claude-3-5-sonnet-20241022", "claude-3-haiku-20240307", "gpt-4o", "gpt-4o-mini"]:
            if m not in model_options:
                model_options.append(m)
        st.selectbox("Model", model_options, key="selected_model")

        st.slider(
            "Max Response Tokens",
            min_value=256, max_value=8000,
            value=chat_manager.max_tokens,
            step=256, key="chat_max_tokens",
        )

        st.divider()

        st.text_area(
            "Custom System Prompt (overrides default if set)",
            key="custom_system_prompt",
            height=120,
            placeholder="Leave empty to use the default assistant prompt with optional file context.",
        )


# ── Entry point ───────────────────────────────────────────────────────────────

def render_chat():
    """Top-level chat view — switches between Agent Mode and normal chat."""
    if st.session_state.get("agent_mode", False):
        st.subheader("Coding Agent")
        render_agent_mode()
    else:
        st.subheader("Chat with AI Assistant")
        render_normal_chat()


if __name__ == "__main__":
    render_chat()