import asyncio
import contextlib
import os
import re
import sys
import threading
import time
from collections import deque
from pathlib import Path
from typing import Any, Optional

import yaml
from fastmcp import FastMCP

from . import co_datascientist_api
from .settings import settings
from .workflow_runner import workflow_runner

mcp = FastMCP("CoDatascientist")

START_MARKER = "# CO_DATASCIENTIST_BLOCK_START"
END_MARKER = "# CO_DATASCIENTIST_BLOCK_END"
IGNORE_DIRS = {"results", "current_runs", "co_datascientist_checkpoints", "co_datascientist_runs", "__pycache__", ".git", ".venv", "venv"}
DEFAULT_MCP_CONFIG_NAME = "config.tropiflo.yaml"

# Global append-only event log and per-client cursors.
_event_log: deque[tuple[int, str]] = deque(maxlen=5000)
_event_lock = threading.Lock()
_next_event_id = 1

_stream_state_lock = threading.Lock()
_stream_states: dict[str, dict[str, float | int]] = {}

_workflow_start_lock = threading.Lock()
_pipeline_state_lock = threading.Lock()
_pipeline_sessions: dict[str, dict[str, Any]] = {}

_MIN_STREAM_POLL_SECONDS = float(os.getenv("CO_DATASCIENTIST_STREAM_MIN_POLL_SECONDS", "30"))


def _normalize_client_id(client_id: Optional[str]) -> str:
    if client_id and client_id.strip():
        return client_id.strip()
    return "default"


def _append_event(message: str) -> None:
    global _next_event_id
    if not message.strip():
        return
    with _event_lock:
        _event_log.append((_next_event_id, message))
        _next_event_id += 1


def _get_events_since(last_event_id: int) -> tuple[list[str], int]:
    with _event_lock:
        new_events = [msg for event_id, msg in _event_log if event_id > last_event_id]
        latest_event_id = _event_log[-1][0] if _event_log else last_event_id
    return new_events, latest_event_id


@contextlib.contextmanager
def _tee_output():
    """Tee stdout/stderr into event log while still printing to console."""

    class Tee:
        def __init__(self, original):
            self.original = original
            self._buf = ""

        def write(self, data):
            self.original.write(data)
            self._buf += data
            while "\n" in self._buf:
                line, self._buf = self._buf.split("\n", 1)
                if line.strip():
                    _append_event(line)

        def flush(self):
            self.original.flush()

    orig_out, orig_err = sys.stdout, sys.stderr
    tee_out, tee_err = Tee(orig_out), Tee(orig_err)
    try:
        sys.stdout, sys.stderr = tee_out, tee_err
        yield
    finally:
        sys.stdout, sys.stderr = orig_out, orig_err
        if getattr(tee_out, "_buf", "").strip():
            _append_event(tee_out._buf.strip())
        if getattr(tee_err, "_buf", "").strip():
            _append_event(tee_err._buf.strip())


def _load_config(config_path: Optional[str]) -> dict:
    """Load YAML config and apply settings overrides."""
    if not config_path:
        return {}
    config_file = Path(config_path)
    if not config_file.exists():
        raise FileNotFoundError(f"Config file not found: {config_path}")

    with config_file.open("r", encoding="utf-8") as fh:
        config = yaml.safe_load(fh) or {}

    if "api_key" in config:
        settings.set_api_key_from_config(config["api_key"])

    backend_url = config.get("backend_url") or config.get("co_datascientist_backend_url")
    backend_url_dev = config.get("backend_url_dev") or config.get("co_datascientist_backend_url_dev")
    if backend_url or backend_url_dev:
        settings.set_backend_urls_from_config(backend_url, backend_url_dev)

    return config


def _infer_entry_command(working_directory: Path) -> Optional[str]:
    candidates = [
        ("run.py", "python run.py"),
        ("main.py", "python main.py"),
        ("train.py", "python train.py"),
        ("run.sh", "bash run.sh"),
    ]
    for filename, command in candidates:
        if (working_directory / filename).exists():
            return command

    py_files = [
        p.name
        for p in working_directory.iterdir()
        if p.is_file() and p.suffix == ".py" and not p.name.startswith(".")
    ]
    if len(py_files) == 1:
        return f"python {py_files[0]}"
    return None


def _ensure_entry_command(config: dict, working_directory: Path, entry_command: Optional[str]) -> None:
    if entry_command:
        config["entry_command"] = entry_command
    if config.get("entry_command"):
        return

    inferred = _infer_entry_command(working_directory)
    if inferred:
        config["entry_command"] = inferred
        return

    raise ValueError(
        "No entry_command provided and could not auto-detect one. "
        "Provide entry_command (e.g., 'python your_script.py') or add run.py/main.py/train.py."
    )


def _normalize_runtime_config(config: dict, use_cached_qa: bool) -> dict:
    config.setdefault("mode", "local")
    try:
        config["parallel"] = max(1, int(config.get("parallel", 1)))
    except Exception:
        config["parallel"] = 1
    # Default to non-interactive MCP behavior.
    config["use_cached_qa"] = True if use_cached_qa else bool(config.get("use_cached_qa", True))
    return config


def _ensure_api_key() -> Optional[str]:
    try:
        settings.get_api_key()
        if not settings.api_key or not settings.api_key.get_secret_value():
            return "No API key found. Please run 'set-token' or add api_key to your config."
    except Exception as exc:
        return f"Error loading API key: {exc}"
    return None


async def _usage_startup_message() -> str:
    try:
        usage_status = await co_datascientist_api.get_user_usage_status()
        msg = f"Usage: ${usage_status['current_usage_usd']:.2f}/${usage_status['limit_usd']:.2f} ({usage_status['usage_percentage']:.1f}%)"
        if usage_status["is_blocked"]:
            return f"BLOCKED: Usage limit exceeded. {msg}"
        return msg
    except Exception as exc:
        return f"Usage status unavailable ({exc})"


async def _ensure_usage_allows_start() -> tuple[Optional[str], str]:
    """
    Validate backend usage/auth before mutating files or starting workflow.
    Returns (error, usage_message).
    """
    try:
        usage_status = await co_datascientist_api.get_user_usage_status()
    except Exception as exc:
        msg = str(exc)
        if "Invalid token" in msg or "401" in msg or "Unauthorized" in msg:
            return ("Invalid token. Please refresh your Tropiflo API key before running optimization.", "")
        return (f"Could not verify usage status: {exc}", "")

    usage_msg = (
        f"Usage: ${usage_status['current_usage_usd']:.2f}/${usage_status['limit_usd']:.2f} "
        f"({usage_status['usage_percentage']:.1f}%)"
    )
    if usage_status["is_blocked"]:
        return (f"Usage limit exceeded. {usage_msg}", usage_msg)
    return (None, usage_msg)


async def _run_with_tee(project_path: Path, config: dict):
    with _tee_output():
        await workflow_runner.run_workflow(str(project_path), config, None)


def _workflow_thread_main(project_path: Path, config: dict):
    try:
        asyncio.run(_run_with_tee(project_path, config))
    except Exception as exc:
        _append_event(f"Workflow crashed before completion: {exc}")


def _start_workflow_thread(project_path: Path, config: dict):
    threading.Thread(
        target=lambda: _workflow_thread_main(project_path, config),
        daemon=True,
    ).start()


def _status_snapshot() -> dict:
    duration_seconds = time.time() - workflow_runner.start_timestamp if getattr(workflow_runner, "start_timestamp", 0) else 0
    if workflow_runner.workflow is None:
        return {"status": "not started", "duration_seconds": duration_seconds}
    return {
        "status": workflow_runner.workflow.status_text,
        "info": workflow_runner.workflow.info,
        "finished": workflow_runner.workflow.finished,
        "duration_seconds": duration_seconds,
        "workflow_id": getattr(workflow_runner.workflow, "workflow_id", None),
    }


def _extract_float(text: str) -> Optional[float]:
    match = re.search(r"([+-]?(?:\d+\.?\d*|\.\d+)(?:[eE][+-]?\d+)?)", text)
    if not match:
        return None
    try:
        return float(match.group(1))
    except ValueError:
        return None


def _phase_from_status(status_text: str, finished: bool) -> str:
    if finished:
        return "finished"
    s = (status_text or "").lower()
    if "preflight" in s:
        return "preflight"
    if "baseline" in s:
        return "baseline"
    if "hypoth" in s or "running" in s or "batch" in s:
        return "optimization"
    if "stopped" in s:
        return "stopped"
    return "running"


def _ensure_journal(state: dict[str, Any]) -> dict[str, Any]:
    journal = state.get("journal")
    if journal is None:
        journal = {
            "run_name": None,
            "baseline_kpi": None,
            "best_kpi": None,
            "best_hypothesis": None,
            "hypotheses": [],
            "results": [],
            "pending_hypotheses": [],
            "in_hypothesis_list": False,
            "current_result_hypothesis": None,
            "last_event_ts": None,
            "last_event_text": None,
        }
        state["journal"] = journal
    return journal


def _update_journal_from_events(journal: dict[str, Any], events: list[str]) -> None:
    now = time.time()
    for raw_line in events:
        line = raw_line.strip()
        if not line:
            continue
        journal["last_event_ts"] = now
        journal["last_event_text"] = line

        if line.startswith("Run name:") or line.startswith("🎯 Run name:"):
            run_name = line.split(":", 1)[1].strip() if ":" in line else None
            if run_name:
                journal["run_name"] = run_name
            continue

        if "Generated Hypotheses:" in line:
            journal["in_hypothesis_list"] = True
            journal["pending_hypotheses"] = []
            continue

        if journal.get("in_hypothesis_list"):
            numbered = re.match(r"^\d+\.\s+(.*)$", line)
            if numbered:
                hyp_text = numbered.group(1).strip()
                journal["pending_hypotheses"].append(hyp_text)
                journal["hypotheses"].append({"text": hyp_text, "ts": now})
                continue
            if line.startswith("Testing ") or line.startswith("🏗️") or line.startswith("-" * 10):
                journal["in_hypothesis_list"] = False

        if line.startswith("Baseline KPI:"):
            kpi = _extract_float(line)
            if kpi is not None:
                journal["baseline_kpi"] = kpi
                if journal.get("best_kpi") is None or kpi > journal["best_kpi"]:
                    journal["best_kpi"] = kpi
                    journal["best_hypothesis"] = "baseline"
            continue

        if line.startswith("Current best KPI:"):
            kpi = _extract_float(line)
            if kpi is not None:
                journal["best_kpi"] = kpi
            continue

        if line.startswith("Hypothesis:"):
            journal["current_result_hypothesis"] = line.split(":", 1)[1].strip()
            continue

        if any(tag in line for tag in ("[IMPROVED]", "[WORSE]", "[RESULT]", "[FAILED]", "[DEBUG]")):
            hypothesis = journal.get("current_result_hypothesis")
            outcome = "result"
            if "[IMPROVED]" in line:
                outcome = "improved"
            elif "[WORSE]" in line:
                outcome = "worse"
            elif "[FAILED]" in line:
                outcome = "failed"
            elif "[DEBUG]" in line:
                outcome = "debug"

            kpi = _extract_float(line) if "KPI" in line else None
            entry = {
                "hypothesis": hypothesis,
                "outcome": outcome,
                "kpi": kpi,
                "ts": now,
            }
            journal["results"].append(entry)
            if outcome == "improved" and hypothesis:
                journal["best_hypothesis"] = hypothesis
            if kpi is not None:
                if journal.get("best_kpi") is None or kpi > journal["best_kpi"]:
                    journal["best_kpi"] = kpi
                    if hypothesis:
                        journal["best_hypothesis"] = hypothesis
            continue


def _build_progress_summary(
    status: dict[str, Any],
    journal: dict[str, Any],
    throttled: bool,
    retry_after_seconds: Optional[int],
) -> tuple[dict[str, Any], str]:
    status_text = str(status.get("status", ""))
    finished = bool(status.get("finished", False))
    phase = _phase_from_status(status_text, finished)
    last_event_ts = journal.get("last_event_ts")
    last_event_age = None
    if isinstance(last_event_ts, (int, float)):
        last_event_age = int(max(0, time.time() - float(last_event_ts)))

    recent_hypotheses = [h.get("text") for h in journal.get("hypotheses", [])[-3:] if h.get("text")]
    recent_results = journal.get("results", [])[-3:]

    summary = {
        "phase": phase,
        "run_name": journal.get("run_name"),
        "status_text": status_text,
        "baseline_kpi": journal.get("baseline_kpi"),
        "best_kpi": journal.get("best_kpi"),
        "best_hypothesis": journal.get("best_hypothesis"),
        "recent_hypotheses": recent_hypotheses,
        "recent_results": recent_results,
        "last_event_seconds_ago": last_event_age,
        "total_hypotheses_generated": len(journal.get("hypotheses", [])),
        "total_results_recorded": len(journal.get("results", [])),
    }

    if finished:
        update = "Optimization finished. "
    elif phase == "preflight":
        update = "Tropiflo is preparing context and preflight checks. "
    elif phase == "baseline":
        update = "Tropiflo is running or validating the baseline execution. "
    elif phase == "optimization":
        update = "Tropiflo is generating/testing hypotheses and scoring results. "
    else:
        update = "Tropiflo is running. "

    if summary["best_kpi"] is not None:
        update += f"Best KPI so far: {summary['best_kpi']}. "
    if recent_results:
        last = recent_results[-1]
        hyp = last.get("hypothesis") or "latest hypothesis"
        outcome = last.get("outcome")
        kpi = last.get("kpi")
        if kpi is not None:
            update += f"Latest result: {hyp} -> {outcome} (KPI {kpi}). "
        else:
            update += f"Latest result: {hyp} -> {outcome}. "
    elif recent_hypotheses:
        update += f"Generated {len(summary['recent_hypotheses'])} recent hypotheses; awaiting execution results. "

    if throttled and retry_after_seconds:
        update += f"Next recommended poll in ~{retry_after_seconds}s."
    elif last_event_age is not None and last_event_age > 30 and not finished:
        update += "No new logs recently; this can be normal while backend generates ideas or jobs are running."

    return summary, update.strip()


def _list_python_files(working_directory: Path) -> list[Path]:
    files: list[Path] = []
    for path in working_directory.rglob("*.py"):
        rel_parts = set(path.relative_to(working_directory).parts)
        if rel_parts & IGNORE_DIRS:
            continue
        files.append(path)
    return sorted(files)


def _has_kpi_print(code: str) -> bool:
    # Accept common forms: print(f"KPI: {x}"), print("KPI:", x), logger text containing "KPI:"
    return bool(re.search(r"KPI\s*[:=]", code, flags=re.IGNORECASE))


def _choose_target_file(working_directory: Path, explicit_target_file: Optional[str]) -> Path:
    if explicit_target_file:
        target = Path(explicit_target_file)
        if not target.is_absolute():
            target = working_directory / target
        return target.resolve()

    preferred = ["train.py", "main.py", "run.py"]
    for name in preferred:
        candidate = working_directory / name
        if candidate.exists():
            return candidate

    py_files = _list_python_files(working_directory)
    if not py_files:
        raise ValueError("No Python files found in the project.")
    return py_files[0]


def _find_first_non_import_line(lines: list[str]) -> int:
    idx = 0
    if lines and lines[0].startswith("#!"):
        idx = 1
    if idx < len(lines) and re.match(r"^#.*coding[:=]\s*[-\w.]+", lines[idx]):
        idx += 1
    while idx < len(lines) and (lines[idx].strip() == "" or lines[idx].strip().startswith("#")):
        idx += 1
    while idx < len(lines):
        stripped = lines[idx].strip()
        if stripped.startswith("import ") or stripped.startswith("from "):
            idx += 1
            continue
        break
    return min(idx, len(lines))


def _insert_markers_if_missing(code: str) -> tuple[str, bool]:
    if START_MARKER in code and END_MARKER in code:
        return code, False

    lines = code.splitlines()
    insert_at = _find_first_non_import_line(lines)
    new_lines = lines[:insert_at] + [START_MARKER] + lines[insert_at:] + [END_MARKER]
    new_code = "\n".join(new_lines)
    if code.endswith("\n"):
        new_code += "\n"
    return new_code, True


def _append_kpi_placeholder_if_missing(code: str) -> tuple[str, bool]:
    if _has_kpi_print(code):
        return code, False
    suffix = (
        "\n\n# TODO: replace this placeholder with your real KPI metric.\n"
        "print(\"KPI:\", 0.0)\n"
    )
    return (code + suffix), True


def _write_default_config(
    working_directory: Path,
    entry_command: str,
    parallel: int,
    config_path: Optional[str],
) -> Path:
    path = Path(config_path).resolve() if config_path else (working_directory / DEFAULT_MCP_CONFIG_NAME).resolve()
    data: dict[str, Any] = {}
    if path.exists():
        with path.open("r", encoding="utf-8") as fh:
            data = yaml.safe_load(fh) or {}
    data.setdefault("mode", "local")
    data["entry_command"] = entry_command
    data["parallel"] = max(1, parallel)
    data.setdefault("use_cached_qa", True)
    with path.open("w", encoding="utf-8") as fh:
        yaml.safe_dump(data, fh, sort_keys=False)
    return path


def _prepare_project_sync(
    working_directory: str,
    target_file: Optional[str] = None,
    entry_command: Optional[str] = None,
    add_kpi_placeholder: bool = False,
    create_config: bool = True,
    config_path: Optional[str] = None,
    parallel: int = 1,
    dry_run: bool = False,
) -> dict:
    project_path = Path(working_directory).resolve()
    if not project_path.exists():
        return {"error": f"Working directory does not exist: {project_path}"}

    try:
        chosen_file = _choose_target_file(project_path, target_file)
    except Exception as exc:
        return {"error": str(exc)}

    if not chosen_file.exists():
        return {"error": f"Target file does not exist: {chosen_file}"}

    original = chosen_file.read_text(encoding="utf-8")
    updated, markers_added = _insert_markers_if_missing(original)
    kpi_added = False
    if add_kpi_placeholder:
        updated, kpi_added = _append_kpi_placeholder_if_missing(updated)

    changed = updated != original
    if changed and not dry_run:
        chosen_file.write_text(updated, encoding="utf-8")

    resolved_entry = entry_command or _infer_entry_command(project_path)
    if not resolved_entry:
        resolved_entry = f"python {chosen_file.name}"

    resolved_config_path = None
    if create_config and not dry_run:
        resolved_config_path = _write_default_config(
            working_directory=project_path,
            entry_command=resolved_entry,
            parallel=parallel,
            config_path=config_path,
        )

    return {
        "working_directory": str(project_path),
        "target_file": str(chosen_file),
        "changed": changed,
        "markers_added": markers_added,
        "kpi_placeholder_added": kpi_added,
        "dry_run": dry_run,
        "entry_command": resolved_entry,
        "config_path": str(resolved_config_path) if resolved_config_path else (str(config_path) if config_path else None),
        "next_step": "Call optimize_and_stream with the returned entry_command/config_path and a stable client_id.",
    }


@mcp.tool()
async def inspect_project_for_tropiflo(
    working_directory: str = ".",
    entry_command: Optional[str] = None,
) -> dict:
    """
    Inspect a project and report whether it is ready for Tropiflo optimization.
    """
    project_path = Path(working_directory).resolve()
    if not project_path.exists():
        return {"error": f"Working directory does not exist: {project_path}"}

    py_files = _list_python_files(project_path)
    files_with_blocks = []
    files_with_kpi = []
    for file_path in py_files:
        content = file_path.read_text(encoding="utf-8")
        rel = str(file_path.relative_to(project_path))
        if START_MARKER in content and END_MARKER in content:
            files_with_blocks.append(rel)
        if _has_kpi_print(content):
            files_with_kpi.append(rel)

    inferred_entry = entry_command or _infer_entry_command(project_path)
    ready = bool(files_with_blocks) and bool(inferred_entry)

    next_steps = []
    if not files_with_blocks:
        next_steps.append("Run setup_project_for_tropiflo to auto-insert CO_DATASCIENTIST block markers.")
    if not files_with_kpi:
        next_steps.append("Add a KPI print line like: print(f'KPI: {metric:.4f}') or enable kpi placeholder during setup.")
    if not inferred_entry:
        next_steps.append("Provide an explicit entry_command, e.g. 'python train.py'.")
    if not next_steps:
        next_steps.append("Project looks ready. Start with optimize_and_stream.")

    return {
        "working_directory": str(project_path),
        "python_files_count": len(py_files),
        "files_with_blocks": files_with_blocks,
        "files_with_kpi_print": files_with_kpi,
        "suggested_entry_command": inferred_entry,
        "ready_to_optimize": ready,
        "next_steps": next_steps,
    }


@mcp.tool()
async def setup_project_for_tropiflo(
    working_directory: str = ".",
    target_file: Optional[str] = None,
    entry_command: Optional[str] = None,
    add_kpi_placeholder: bool = False,
    create_config: bool = True,
    config_path: Optional[str] = None,
    parallel: int = 1,
    dry_run: bool = False,
) -> dict:
    """
    Auto-prepare a project for Tropiflo by adding evolution markers and creating config.
    """
    return _prepare_project_sync(
        working_directory=working_directory,
        target_file=target_file,
        entry_command=entry_command,
        add_kpi_placeholder=add_kpi_placeholder,
        create_config=create_config,
        config_path=config_path,
        parallel=parallel,
        dry_run=dry_run,
    )


@mcp.tool()
async def optimize_pipeline(
    working_directory: str = ".",
    session_id: Optional[str] = None,
    entry_command: Optional[str] = None,
    target_file: Optional[str] = None,
    add_kpi_placeholder: bool = False,
    parallel: int = 1,
    stop: bool = False,
    reset_session: bool = False,
) -> dict:
    """
    Single-tool orchestration for chat agents:
      - first call: auto-setup project + start Tropiflo
      - subsequent calls: poll progress/events
      - optional stop/reset controls

    Recommended usage:
      1) Call with a stable session_id.
      2) Keep calling until finished=true.
      3) Display events_text each call.
    """
    sid = _normalize_client_id(session_id)
    project_path = Path(working_directory).resolve()
    if not project_path.exists():
        return {"error": f"Working directory does not exist: {project_path}"}

    if reset_session:
        with _pipeline_state_lock:
            _pipeline_sessions.pop(sid, None)

    if stop:
        stop_msg = await stop_workflow()
        poll = await optimize_and_stream(
            working_directory=str(project_path),
            client_id=sid,
            use_cached_qa=True,
            auto_setup=False,
        )
        poll["session_id"] = sid
        poll["orchestration"] = {"stop_requested": True, "message": stop_msg}
        return poll

    setup_info = None
    with _pipeline_state_lock:
        state = _pipeline_sessions.get(sid)

    if not state:
        api_key_error = _ensure_api_key()
        if api_key_error:
            return {"error": api_key_error}
        usage_error, _ = await _ensure_usage_allows_start()
        if usage_error:
            return {"error": usage_error}

        setup_info = _prepare_project_sync(
            working_directory=str(project_path),
            target_file=target_file,
            entry_command=entry_command,
            add_kpi_placeholder=add_kpi_placeholder,
            create_config=True,
            config_path=None,
            parallel=max(1, parallel),
            dry_run=False,
        )
        if setup_info.get("error"):
            return {"error": f"Setup failed: {setup_info['error']}"}

        state = {
            "working_directory": str(project_path),
            "entry_command": setup_info.get("entry_command"),
            "config_path": setup_info.get("config_path"),
            "created_at": time.time(),
            "journal": None,
        }
        with _pipeline_state_lock:
            _pipeline_sessions[sid] = state

    response = await optimize_and_stream(
        working_directory=str(project_path),
        client_id=sid,
        entry_command=entry_command or state.get("entry_command"),
        config_path=state.get("config_path"),
        use_cached_qa=True,
        auto_setup=False,
    )
    with _pipeline_state_lock:
        active_state = _pipeline_sessions.setdefault(sid, state)
        journal = _ensure_journal(active_state)
    _update_journal_from_events(journal, response.get("events") or [])

    throttled = bool(response.get("throttled", False))
    retry_after = response.get("retry_after_seconds")
    retry_after_int = int(retry_after) if isinstance(retry_after, (int, float)) else None
    progress_summary, user_update = _build_progress_summary(
        status=response.get("status") or {},
        journal=journal,
        throttled=throttled,
        retry_after_seconds=retry_after_int,
    )

    response["session_id"] = sid
    response["orchestration"] = {
        "single_tool_mode": True,
        "setup_performed": bool(setup_info),
        "next_action": "Call optimize_pipeline again with the same session_id until finished=true.",
    }
    response["progress_summary"] = progress_summary
    response["user_update"] = user_update
    response["recommended_poll_seconds"] = (
        max(int(_MIN_STREAM_POLL_SECONDS), retry_after_int or int(_MIN_STREAM_POLL_SECONDS))
    )
    if setup_info:
        response["setup"] = setup_info
    return response


@mcp.tool()
async def stop_workflow() -> str:
    """
    Request stopping the current workflow.
    """
    if workflow_runner.workflow is None:
        return "No workflow is currently running."
    _append_event("Stop requested by client.")
    workflow_runner.should_stop_workflow = True
    try:
        if getattr(workflow_runner, "workflow", None) and getattr(workflow_runner.workflow, "workflow_id", None):
            await co_datascientist_api.stop_workflow(workflow_runner.workflow.workflow_id)
    except Exception:
        pass
    return "Workflow scheduled to stop. Keep polling optimize_and_stream until finished=true."


@mcp.tool()
async def optimize_and_stream(
    working_directory: str = ".",
    client_id: Optional[str] = None,
    entry_command: Optional[str] = None,
    config_path: Optional[str] = None,
    use_cached_qa: bool = True,
    auto_setup: bool = False,
    setup_target_file: Optional[str] = None,
    add_kpi_placeholder: bool = False,
) -> dict:
    """
    Start or continue optimization and stream logs/status.

    Recommended usage:
      1) Call with a stable client_id (e.g. your chat session id).
      2) Poll this tool repeatedly until finished=true.
      3) Respect throttled/retry_after_seconds when provided.
    """
    started = False
    setup_info = None
    cid = _normalize_client_id(client_id)
    project_path = Path(working_directory).resolve()
    if not project_path.exists():
        return {"error": f"Working directory does not exist: {project_path}"}

    api_key_error = _ensure_api_key()
    if api_key_error:
        return {"error": api_key_error}
    usage_error, usage_msg = await _ensure_usage_allows_start()
    if usage_error:
        return {"error": usage_error}

    if auto_setup:
        setup_info = _prepare_project_sync(
            working_directory=str(project_path),
            target_file=setup_target_file,
            entry_command=entry_command,
            add_kpi_placeholder=add_kpi_placeholder,
            create_config=not bool(config_path),
            config_path=config_path,
            parallel=1,
            dry_run=False,
        )
        if setup_info.get("error"):
            return {"error": f"Auto-setup failed: {setup_info['error']}"}
        if not config_path and setup_info.get("config_path"):
            config_path = setup_info["config_path"]
        if not entry_command and setup_info.get("entry_command"):
            entry_command = setup_info["entry_command"]

    should_start = workflow_runner.workflow is None or getattr(workflow_runner.workflow, "finished", False)
    if should_start:
        with _workflow_start_lock:
            should_start = workflow_runner.workflow is None or getattr(workflow_runner.workflow, "finished", False)
            if should_start:
                try:
                    config = _load_config(config_path)
                    _ensure_entry_command(config, project_path, entry_command)
                    config = _normalize_runtime_config(config, use_cached_qa=use_cached_qa)
                except Exception as exc:
                    return {"error": f"Failed preparing workflow config: {exc}"}

                _start_workflow_thread(project_path, config)
                started = True
                _append_event(
                    f"Started workflow at {project_path} with entry_command='{config.get('entry_command')}' "
                    f"parallel={config.get('parallel')} engine={config.get('engine', 'EVOLVE_HYPOTHESIS')}. {usage_msg}"
                )

    now = time.time()
    with _stream_state_lock:
        state = _stream_states.setdefault(cid, {"last_event_id": 0, "last_poll_ts": 0.0})
        last_event_id = int(state["last_event_id"])
        last_poll_ts = float(state["last_poll_ts"])

    events, latest_event_id = _get_events_since(last_event_id)
    throttled = last_poll_ts and (now - last_poll_ts) < _MIN_STREAM_POLL_SECONDS

    with _stream_state_lock:
        state = _stream_states.setdefault(cid, {"last_event_id": 0, "last_poll_ts": 0.0})
        state["last_event_id"] = latest_event_id
        state["last_poll_ts"] = now

    status = _status_snapshot()
    response: dict[str, Any] = {
        "started": started,
        "status": status,
        "events": events,
        "events_text": "\n".join(events) if events else "",
        "finished": status.get("finished", False),
        "client_id": cid,
        "message": "Events included" if events else "No new events yet; still running",
    }
    if setup_info:
        response["setup"] = setup_info
    if throttled:
        retry_after = max(1, int(_MIN_STREAM_POLL_SECONDS - (now - last_poll_ts)))
        response["throttled"] = True
        response["retry_after_seconds"] = retry_after
        response["message"] = "Throttled: please wait before the next poll."
    return response


@mcp.tool()
async def tropiflo_this(
    working_directory: str = ".",
    client_id: Optional[str] = None,
    entry_command: Optional[str] = None,
    target_file: Optional[str] = None,
    add_kpi_placeholder: bool = False,
) -> dict:
    """
    One-shot helper for chat agents:
      - auto-prepare project markers/config
      - start optimization
      - return first status/events payload
    """
    return await optimize_and_stream(
        working_directory=working_directory,
        client_id=client_id,
        entry_command=entry_command,
        config_path=None,
        use_cached_qa=True,
        auto_setup=True,
        setup_target_file=target_file,
        add_kpi_placeholder=add_kpi_placeholder,
    )


@mcp.tool()
async def get_usage_status() -> dict:
    """
    Get your current usage status and remaining balance. Shows a quick overview similar to 'co-datascientist status' command.
    
    returns:
        dictionary with usage information including:
        - current usage vs limit
        - remaining balance  
        - usage percentage
        - status indicator
        - helpful tips
    """
    try:
        usage_status = await co_datascientist_api.get_user_usage_status()
        
        # Determine status message and emoji
        percentage = usage_status['usage_percentage']
        if usage_status['is_blocked']:
            status_msg = "🚨 BLOCKED - Free tokens exhausted! Contact support or wait for reset."
        elif percentage >= 90:
            status_msg = f"🟥 CRITICAL - Only ${usage_status['remaining_usd']:.2f} remaining!"
        elif percentage >= 80:
            status_msg = f"🟨 WARNING - Approaching limit ({percentage:.1f}% used)"
        elif percentage >= 50:
            status_msg = f"🟦 MODERATE - {percentage:.1f}% of limit used"
        else:
            status_msg = f"🟩 GOOD - Plenty of free tokens remaining"
        
        # Create progress bar representation
        bar_width = 20
        filled = int(bar_width * percentage / 100)
        bar = "█" * filled + "░" * (bar_width - filled)
        
        return {
            "summary": f"💰 Quick Usage Status",
            "used": f"${usage_status['current_usage_usd']:.2f} / ${usage_status['limit_usd']:.2f}",
            "remaining": f"${usage_status['remaining_usd']:.2f}",
            "progress_bar": f"[{bar}] {percentage:.1f}%",
            "status": status_msg,
            "is_blocked": usage_status['is_blocked'],
            "tips": [
                "Use 'get_cost_summary' for basic cost breakdown",
                "Use 'get_detailed_costs' for full analysis",
                "Monitor live progress with 'optimize_and_stream'"
            ]
        }
        
    except Exception as e:
        return {
            "error": f"Error getting usage status: {e}",
            "tips": ["Check your network connection", "Verify MCP server is running correctly"]
        }


@mcp.tool()
async def get_cost_summary() -> dict:
    """
    Get a summary of your usage costs and token consumption. Similar to 'co-datascientist costs' command.
    
    returns:
        dictionary with cost summary including:
        - total cost
        - usage limits
        - token counts
        - workflow counts
    """
    try:
        costs_response = await co_datascientist_api.get_user_costs_summary()
        usage_status = await co_datascientist_api.get_user_usage_status()
        
        # Status indicator
        if usage_status['is_blocked']:
            status_indicator = "🚨 BLOCKED - Free tokens exhausted!"
            status_detail = f"You've used ${usage_status['current_usage_usd']:.2f} of your ${usage_status['limit_usd']:.2f} limit."
        elif usage_status['usage_percentage'] >= 80:
            status_indicator = f"⚠️  Approaching limit - {usage_status['usage_percentage']:.1f}% used"
            status_detail = f"${usage_status['remaining_usd']:.2f} remaining"
        else:
            status_indicator = f"✅ Active - {usage_status['usage_percentage']:.1f}% of limit used"
            status_detail = f"${usage_status['remaining_usd']:.2f} remaining"
        
        return {
            "title": "💰 Co-DataScientist Usage Summary",
            "total_cost": f"${costs_response['total_cost_usd']:.8f}",
            "usage_limit": f"${usage_status['limit_usd']:.2f}",
            "remaining": f"${usage_status['remaining_usd']:.2f} ({usage_status['usage_percentage']:.1f}% used)",
            "status": status_indicator,
            "status_detail": status_detail,
            "total_tokens": f"{costs_response['total_tokens']:,}",
            "workflows_completed": costs_response['workflows_completed'],
            "last_updated": costs_response.get('last_updated', 'Unknown'),
            "tip": "💡 Use 'get_detailed_costs' for full breakdown"
        }
        
    except Exception as e:
        return {
            "error": f"Error getting cost summary: {e}",
            "tip": "Check your connection and try again"
        }


@mcp.tool()
async def get_detailed_costs() -> dict:
    """
    Get detailed cost breakdown including all workflows and model calls. Similar to 'co-datascientist costs --detailed' command.
    
    returns:
        dictionary with detailed cost information including:
        - total costs and limits
        - per-workflow breakdown
        - model call details
        - usage status
    """
    try:
        costs_response = await co_datascientist_api.get_user_costs()
        usage_status = await co_datascientist_api.get_user_usage_status()
        
        # Status with emoji
        if usage_status['is_blocked']:
            status_msg = f"🚨 Status: BLOCKED (limit exceeded)"
        elif usage_status['usage_percentage'] >= 80:
            status_msg = f"⚠️  Status: Approaching limit ({usage_status['usage_percentage']:.1f}%)"
        else:
            status_msg = f"✅ Status: Active ({usage_status['usage_percentage']:.1f}% used)"
        
        # Build workflow breakdown
        workflow_breakdown = []
        if costs_response['workflows']:
            for workflow_id, workflow_data in costs_response['workflows'].items():
                workflow_info = {
                    "id": workflow_id[:8] + "...",
                    "cost": f"${workflow_data['cost']:.8f}",
                    "tokens": f"{workflow_data['input_tokens'] + workflow_data['output_tokens']:,}",
                    "model_calls": len(workflow_data['model_calls'])
                }
                
                # Add recent model calls
                recent_calls = []
                for call in workflow_data['model_calls'][-3:]:  # Last 3 calls
                    recent_calls.append({
                        "model": call['model'],
                        "cost": f"${call['cost']:.8f}",
                        "tokens": f"{call['input_tokens']}+{call['output_tokens']}"
                    })
                workflow_info["recent_calls"] = recent_calls
                
                if len(workflow_data['model_calls']) > 3:
                    workflow_info["additional_calls"] = len(workflow_data['model_calls']) - 3
                
                workflow_breakdown.append(workflow_info)
        
        return {
            "title": "💰 Co-DataScientist Usage Details",
            "total_cost": f"${costs_response['total_cost_usd']:.8f}",
            "usage_limit": f"${usage_status['limit_usd']:.2f}",
            "remaining": f"${usage_status['remaining_usd']:.2f}",
            "usage_percentage": f"{usage_status['usage_percentage']:.1f}%",
            "status": status_msg,
            "total_tokens": f"{costs_response['total_tokens']:,} ({costs_response['total_input_tokens']:,} input + {costs_response['total_output_tokens']:,} output)",
            "workflows_count": costs_response['workflows_count'],
            "last_updated": costs_response.get('last_updated', 'Unknown'),
            "workflow_breakdown": workflow_breakdown
        }
        
    except Exception as e:
        return {
            "error": f"Error getting detailed costs: {e}",
            "tip": "Check your connection and try again"
        }


async def run_mcp_server():
    # Use SSE transport so IDE clients can connect over HTTP
    # Default away from 8000 to avoid clashing with backend; env overrides still respected
    port = int(os.getenv("CO_DATASCIENTIST_MCP_PORT", "8765"))
    await mcp.run_async(
        transport="sse",
        host=settings.host,
        port=port,
    )
