"""
Sequential Thinking MCP Server
===============================
MCP server that provides structured, step-by-step thinking tools for
complex problem solving. Supports creating thinking chains, branching
hypotheses, revising earlier steps, and synthesizing conclusions.
Helps LLMs reason through multi-step problems methodically.
"""

import json
import os
import time
from typing import Optional, List, Dict, Any
from datetime import datetime, timezone
from enum import Enum

from mcp.server.fastmcp import FastMCP

# ---------------------------------------------------------------------------
# MCP Server
# ---------------------------------------------------------------------------

mcp = FastMCP("sequential_thinking_mcp")

# ---------------------------------------------------------------------------
# In-memory session storage
# ---------------------------------------------------------------------------

_sessions: Dict[str, Dict[str, Any]] = {}


class ThoughtType(str, Enum):
    OBSERVATION = "observation"
    HYPOTHESIS = "hypothesis"
    ANALYSIS = "analysis"
    CONCLUSION = "conclusion"
    REVISION = "revision"
    QUESTION = "question"
    EVIDENCE = "evidence"
    COUNTER_ARGUMENT = "counter_argument"


# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------


@mcp.tool()
async def start_thinking_session(
    topic: str,
    context: Optional[str] = None,
    session_id: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Start a new sequential thinking session for a topic.

    Args:
        topic: The main topic or problem to think through
        context: Optional background context or constraints
        session_id: Optional custom session ID (auto-generated if not provided)
    """
    sid = session_id or f"think_{int(time.time() * 1000)}"
    now = datetime.now(timezone.utc).isoformat()

    _sessions[sid] = {
        "session_id": sid,
        "topic": topic,
        "context": context,
        "created_at": now,
        "updated_at": now,
        "thoughts": [],
        "branches": {},
        "status": "active",
        "conclusion": None,
    }

    return {
        "session_id": sid,
        "topic": topic,
        "status": "active",
        "message": "Thinking session started. Add thoughts with add_thought.",
    }


@mcp.tool()
async def add_thought(
    session_id: str,
    thought: str,
    thought_type: str = "analysis",
    confidence: Optional[float] = None,
    references_steps: Optional[List[int]] = None,
    branch: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Add a thought/reasoning step to a thinking session.

    Args:
        session_id: The thinking session ID
        thought: The thought or reasoning content
        thought_type: Type of thought: 'observation', 'hypothesis', 'analysis', 'conclusion', 'revision', 'question', 'evidence', 'counter_argument'
        confidence: Optional confidence level 0.0-1.0
        references_steps: Optional list of step numbers this thought builds on
        branch: Optional branch name for alternative reasoning paths
    """
    if session_id not in _sessions:
        return {"error": f"Session '{session_id}' not found"}

    session = _sessions[session_id]
    now = datetime.now(timezone.utc).isoformat()

    # Determine step number
    if branch:
        if branch not in session["branches"]:
            session["branches"][branch] = []
        step_num = len(session["branches"][branch]) + 1
        target_list = session["branches"][branch]
    else:
        step_num = len(session["thoughts"]) + 1
        target_list = session["thoughts"]

    thought_entry = {
        "step": step_num,
        "thought": thought,
        "type": thought_type,
        "confidence": confidence,
        "references": references_steps or [],
        "branch": branch,
        "timestamp": now,
    }

    target_list.append(thought_entry)
    session["updated_at"] = now

    return {
        "session_id": session_id,
        "step": step_num,
        "branch": branch,
        "type": thought_type,
        "total_steps": len(session["thoughts"]),
        "total_branches": len(session["branches"]),
    }


@mcp.tool()
async def revise_thought(
    session_id: str,
    step_number: int,
    revised_thought: str,
    reason: str,
    branch: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Revise an earlier thought step with new reasoning.

    Args:
        session_id: The thinking session ID
        step_number: The step number to revise
        revised_thought: The new thought content
        reason: Reason for the revision
        branch: Optional branch name if revising a branched thought
    """
    if session_id not in _sessions:
        return {"error": f"Session '{session_id}' not found"}

    session = _sessions[session_id]

    if branch:
        target_list = session["branches"].get(branch, [])
    else:
        target_list = session["thoughts"]

    # Find the step
    for entry in target_list:
        if entry["step"] == step_number:
            entry["original_thought"] = entry["thought"]
            entry["thought"] = revised_thought
            entry["revision_reason"] = reason
            entry["revised_at"] = datetime.now(timezone.utc).isoformat()

            # Also add a revision note to the main chain
            await add_thought(
                session_id=session_id,
                thought=f"[Revision of step {step_number}] {reason}: {revised_thought}",
                thought_type="revision",
                references_steps=[step_number],
                branch=branch,
            )

            return {
                "status": "revised",
                "step": step_number,
                "reason": reason,
            }

    return {"error": f"Step {step_number} not found"}


@mcp.tool()
async def get_thinking_chain(
    session_id: str,
    branch: Optional[str] = None,
    include_revisions: bool = True,
) -> Dict[str, Any]:
    """
    Get the full chain of thoughts for a session.

    Args:
        session_id: The thinking session ID
        branch: Optional branch name to get (None for main chain)
        include_revisions: Whether to include revision history
    """
    if session_id not in _sessions:
        return {"error": f"Session '{session_id}' not found"}

    session = _sessions[session_id]

    if branch:
        thoughts = session["branches"].get(branch, [])
    else:
        thoughts = session["thoughts"]

    if not include_revisions:
        thoughts = [
            {k: v for k, v in t.items() if k not in ("original_thought", "revision_reason", "revised_at")}
            for t in thoughts
        ]

    return {
        "session_id": session_id,
        "topic": session["topic"],
        "context": session["context"],
        "branch": branch,
        "thoughts": thoughts,
        "total_steps": len(thoughts),
        "branches_available": list(session["branches"].keys()),
        "status": session["status"],
        "conclusion": session["conclusion"],
    }


@mcp.tool()
async def synthesize_conclusion(
    session_id: str,
    conclusion: str,
    confidence: Optional[float] = None,
    key_insights: Optional[List[str]] = None,
) -> Dict[str, Any]:
    """
    Synthesize a conclusion from the thinking chain.

    Args:
        session_id: The thinking session ID
        conclusion: The synthesized conclusion
        confidence: Overall confidence level 0.0-1.0
        key_insights: Optional list of key insights from the thinking process
    """
    if session_id not in _sessions:
        return {"error": f"Session '{session_id}' not found"}

    session = _sessions[session_id]
    now = datetime.now(timezone.utc).isoformat()

    session["conclusion"] = {
        "text": conclusion,
        "confidence": confidence,
        "key_insights": key_insights or [],
        "concluded_at": now,
        "based_on_steps": len(session["thoughts"]),
        "branches_considered": list(session["branches"].keys()),
    }
    session["status"] = "concluded"
    session["updated_at"] = now

    # Add conclusion as final thought
    await add_thought(
        session_id=session_id,
        thought=conclusion,
        thought_type="conclusion",
        confidence=confidence,
    )

    return {
        "session_id": session_id,
        "status": "concluded",
        "conclusion": session["conclusion"],
    }


@mcp.tool()
async def compare_branches(
    session_id: str,
    branch_names: Optional[List[str]] = None,
) -> Dict[str, Any]:
    """
    Compare different reasoning branches in a session.

    Args:
        session_id: The thinking session ID
        branch_names: Optional list of branches to compare (all if not specified)
    """
    if session_id not in _sessions:
        return {"error": f"Session '{session_id}' not found"}

    session = _sessions[session_id]
    branches = branch_names or list(session["branches"].keys())

    comparison = {
        "main_chain": {
            "steps": len(session["thoughts"]),
            "types": _count_types(session["thoughts"]),
            "avg_confidence": _avg_confidence(session["thoughts"]),
        }
    }

    for branch in branches:
        if branch in session["branches"]:
            thoughts = session["branches"][branch]
            comparison[branch] = {
                "steps": len(thoughts),
                "types": _count_types(thoughts),
                "avg_confidence": _avg_confidence(thoughts),
            }

    return {
        "session_id": session_id,
        "comparison": comparison,
    }


def _count_types(thoughts: List[Dict]) -> Dict[str, int]:
    counts: Dict[str, int] = {}
    for t in thoughts:
        tt = t.get("type", "unknown")
        counts[tt] = counts.get(tt, 0) + 1
    return counts


def _avg_confidence(thoughts: List[Dict]) -> Optional[float]:
    confs = [t["confidence"] for t in thoughts if t.get("confidence") is not None]
    if not confs:
        return None
    return round(sum(confs) / len(confs), 3)


@mcp.tool()
async def list_sessions(
    status: Optional[str] = None,
    limit: int = 20,
) -> Dict[str, Any]:
    """
    List all thinking sessions.

    Args:
        status: Optional filter by status: 'active' or 'concluded'
        limit: Maximum sessions to return
    """
    sessions = []
    for sid, session in _sessions.items():
        if status and session["status"] != status:
            continue
        sessions.append({
            "session_id": sid,
            "topic": session["topic"],
            "status": session["status"],
            "total_steps": len(session["thoughts"]),
            "branches": len(session["branches"]),
            "created_at": session["created_at"],
            "updated_at": session["updated_at"],
        })

    sessions.sort(key=lambda x: x["updated_at"], reverse=True)

    return {
        "total": len(sessions),
        "sessions": sessions[:limit],
    }