"""
Sequential Thinking — Gradio MCP App

Ports the Sequential Thinking MCP Server (TypeScript) to a Python Gradio app
that exposes the same tool via Gradio's built-in MCP server support.

Launch:  python app.py
MCP SSE endpoint:  http://localhost:7860/gradio_api/mcp/sse
"""

import json
import os
import sys
from typing import Optional
from uuid import uuid4

# On Windows the default console encoding (cp1252) cannot represent the emoji
# characters that Gradio prints on startup.  Reconfigure stdout/stderr to
# UTF-8 early so those prints succeed.  When running via `run.ps1` (or with
# `python -X utf8 app.py`) the streams are already UTF-8 and this is a no-op.
if sys.platform == "win32":
    for _stream in (sys.stdout, sys.stderr):
        try:
            if hasattr(_stream, "reconfigure") and _stream.encoding.lower() != "utf-8":
                _stream.reconfigure(encoding="utf-8", errors="replace")
        except Exception:
            pass

import gradio as gr
from theme import soft_professional_theme
from chat_logger import log_chat, shutdown_logger

# ---------------------------------------------------------------------------
# Core state (mirrors SequentialThinkingServer in lib.ts)
# ---------------------------------------------------------------------------

_thought_history: list[dict] = []
_branches: dict[str, list[dict]] = {}
_session_id: str = uuid4().hex  # Session identifier for logging


# ---------------------------------------------------------------------------
# Tool logic
# ---------------------------------------------------------------------------

def sequential_thinking(
    thought: str,
    thought_number: int,
    total_thoughts: int,
    next_thought_needed: bool,
    is_revision: bool = False,
    revises_thought: Optional[int] = None,
    branch_from_thought: Optional[int] = None,
    branch_id: Optional[str] = None,
    needs_more_thoughts: bool = False,
) -> dict:
    """
    A detailed tool for dynamic and reflective problem-solving through thoughts.
    This tool helps analyze problems through a flexible thinking process that
    can adapt and evolve. Each thought can build on, question, or revise
    previous insights as understanding deepens.

    Use this tool when:
    - Breaking down complex problems into steps
    - Planning and design with room for revision
    - Analysis that might need course correction
    - Problems where the full scope is not clear initially
    - Tasks that need context maintained over multiple steps

    Args:
        thought: Your current thinking step (analytical step, revision, question, hypothesis, etc.)
        thought_number: Current thought number in sequence (starts at 1)
        total_thoughts: Estimated total thoughts needed — can be adjusted up or down
        next_thought_needed: True if more thinking is required, even if at what seemed like the end
        is_revision: True if this thought revises or corrects a previous thought
        revises_thought: Which thought number is being reconsidered (required when is_revision=True)
        branch_from_thought: Thought number this branch splits off from (for alternative paths)
        branch_id: Unique identifier for this branch (e.g. "alternative-approach")
        needs_more_thoughts: True if you realise more thoughts are needed beyond totalThoughts

    Returns:
        dict with thoughtNumber, totalThoughts, nextThoughtNeeded, branches, thoughtHistoryLength
    """
    global _thought_history, _branches

    # Sanitise optional int inputs that arrive as 0 from Gradio Number widgets
    if revises_thought == 0:
        revises_thought = None
    if branch_from_thought == 0:
        branch_from_thought = None

    # Mirror TypeScript behaviour: expand totalThoughts if thoughtNumber exceeds it
    if thought_number > total_thoughts:
        total_thoughts = thought_number

    thought_data: dict = {
        "thought": thought,
        "thoughtNumber": thought_number,
        "totalThoughts": total_thoughts,
        "isRevision": is_revision,
        "revisesThought": revises_thought,
        "branchFromThought": branch_from_thought,
        "branchId": branch_id if branch_id and branch_id.strip() else None,
        "needsMoreThoughts": needs_more_thoughts,
        "nextThoughtNeeded": next_thought_needed,
    }

    _thought_history.append(thought_data)

    if branch_from_thought and branch_id and branch_id.strip():
        _branches.setdefault(branch_id, []).append(thought_data)

    _log_thought(thought_data)

    # Log to local storage
    log_chat(
        session_id=_session_id,
        model_name="sequential-thinking",
        thought=thought,
        thought_number=thought_number,
        total_thoughts=total_thoughts,
        metadata={
            "is_revision": is_revision,
            "revises_thought": revises_thought,
            "branch_from_thought": branch_from_thought,
            "branch_id": branch_id,
            "needs_more_thoughts": needs_more_thoughts,
            "next_thought_needed": next_thought_needed,
        }
    )

    return {
        "thoughtNumber": thought_number,
        "totalThoughts": total_thoughts,
        "nextThoughtNeeded": next_thought_needed,
        "branches": list(_branches.keys()),
        "thoughtHistoryLength": len(_thought_history),
    }


def reset_session() -> str:
    """
    Reset the thought history and all branches, starting a fresh problem-solving session.

    Returns:
        Confirmation message string.
    """
    global _thought_history, _branches, _session_id
    _thought_history = []
    _branches = {}
    _session_id = uuid4().hex  # Generate new session ID
    return "Session reset. Ready for a new problem-solving session."


def get_history() -> list:
    """
    Return the full thought history for the current session.

    Returns:
        List of all recorded thought objects.
    """
    return _thought_history


# ---------------------------------------------------------------------------
# Logging helper (mirrors lib.ts formatThought, without chalk colours)
# ---------------------------------------------------------------------------

def _log_thought(data: dict) -> None:
    if os.environ.get("DISABLE_THOUGHT_LOGGING", "").lower() == "true":
        return

    t_num = data["thoughtNumber"]
    t_total = data["totalThoughts"]
    thought = data["thought"]

    if data.get("isRevision"):
        label = f"[Revision] {t_num}/{t_total} (revising thought {data.get('revisesThought')})"
    elif data.get("branchFromThought"):
        label = (
            f"[Branch] {t_num}/{t_total} "
            f"(from thought {data['branchFromThought']}, ID: {data.get('branchId')})"
        )
    else:
        label = f"[Thought] {t_num}/{t_total}"

    width = max(len(label), len(thought)) + 4
    border = "-" * width
    try:
        print(f"\n+{border}+")
        print(f"| {label.ljust(width - 2)} |")
        print(f"+{border}+")
        print(f"| {thought.ljust(width - 2)} |")
        print(f"+{border}+")
    except UnicodeEncodeError:
        print(f"\n{label}\n  {thought}")


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

DESCRIPTION = """
# 🧠 Sequential Thinking MCP Server

Dynamic, reflective problem-solving through structured thoughts.  
This Gradio app exposes the **Sequential Thinking** tool as an MCP server.

**MCP endpoint (SSE):** `http://localhost:7860/gradio_api/mcp/sse`
"""

# Load custom CSS
with open('styles.css', 'r') as f:
    custom_css = f.read()

with gr.Blocks(title="Sequential Thinking MCP", theme=soft_professional_theme, css=custom_css) as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Tabs():

        # ------------------------------------------------------------------
        # Tab 1 — process a thought
        # ------------------------------------------------------------------
        with gr.Tab("Process Thought"):
            with gr.Row(equal_height=False):
                with gr.Column(scale=1):
                    thought_input = gr.Textbox(
                        label="Thought",
                        placeholder="Enter your current thinking step…",
                        lines=5,
                    )
                    with gr.Row():
                        thought_number_input = gr.Number(
                            label="Thought Number", value=1, minimum=1, precision=0
                        )
                        total_thoughts_input = gr.Number(
                            label="Total Thoughts (estimate)", value=5, minimum=1, precision=0
                        )
                    next_thought_needed_input = gr.Checkbox(
                        label="Next Thought Needed", value=True
                    )

                    with gr.Accordion("Advanced options", open=False):
                        is_revision_input = gr.Checkbox(label="Is Revision", value=False)
                        revises_thought_input = gr.Number(
                            label="Revises Thought #",
                            minimum=0,
                            precision=0,
                            value=0,
                            info="0 = not applicable",
                        )
                        branch_from_thought_input = gr.Number(
                            label="Branch From Thought #",
                            minimum=0,
                            precision=0,
                            value=0,
                            info="0 = not applicable",
                        )
                        branch_id_input = gr.Textbox(
                            label="Branch ID",
                            placeholder="e.g. alternative-approach",
                        )
                        needs_more_thoughts_input = gr.Checkbox(
                            label="Needs More Thoughts", value=False
                        )

                    with gr.Row():
                        submit_btn = gr.Button("Process Thought", variant="primary")
                        reset_btn = gr.Button("Reset Session", variant="secondary")

                with gr.Column(scale=1):
                    output_json = gr.JSON(label="Result")
                    status_box = gr.Textbox(label="Status", interactive=False, visible=True)

            submit_btn.click(
                fn=sequential_thinking,
                inputs=[
                    thought_input,
                    thought_number_input,
                    total_thoughts_input,
                    next_thought_needed_input,
                    is_revision_input,
                    revises_thought_input,
                    branch_from_thought_input,
                    branch_id_input,
                    needs_more_thoughts_input,
                ],
                outputs=output_json,
            )

            reset_btn.click(fn=reset_session, inputs=[], outputs=status_box)

        # ------------------------------------------------------------------
        # Tab 2 — inspect history
        # ------------------------------------------------------------------
        with gr.Tab("Thought History"):
            refresh_btn = gr.Button("Refresh")
            history_output = gr.JSON(label="All Thoughts")
            refresh_btn.click(fn=get_history, inputs=[], outputs=history_output)

        # ------------------------------------------------------------------
        # Tab 3 — quick reference
        # ------------------------------------------------------------------
        with gr.Tab("MCP Reference"):
            gr.Markdown("""
## Connecting via MCP

Add this to your MCP client configuration:

```json
{
  "mcpServers": {
    "sequential-thinking": {
      "url": "http://localhost:7860/gradio_api/mcp/sse"
    }
  }
}
```

## Exposed MCP Tools

| Tool | Description |
|------|-------------|
| `sequential_thinking` | Process a single thought step |
| `reset_session` | Clear history and branches |
| `get_history` | Return the full thought history |

## Parameters for `sequential_thinking`

| Parameter | Type | Required | Description |
|-----------|------|----------|-------------|
| `thought` | string | ✅ | Current thinking step |
| `thought_number` | int | ✅ | Step index (starts at 1) |
| `total_thoughts` | int | ✅ | Estimated total steps |
| `next_thought_needed` | bool | ✅ | More steps required? |
| `is_revision` | bool | | Revises an earlier thought |
| `revises_thought` | int | | Which thought is revised |
| `branch_from_thought` | int | | Branch origin thought |
| `branch_id` | string | | Branch label |
| `needs_more_thoughts` | bool | | Extend beyond estimate |

## Environment Variables

| Variable | Default | Description |
|----------|---------|-------------|
| `DISABLE_THOUGHT_LOGGING` | `false` | Set to `true` to suppress console output |
""")


if __name__ == "__main__":
    try:
        demo.launch(mcp_server=True)
    finally:
        shutdown_logger()