from __future__ import annotations

import asyncio
import contextvars
import json
import logging
import os
import shutil
import time
import uuid
import subprocess
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from fastapi import APIRouter, HTTPException, Header
from pydantic import BaseModel as PydanticBaseModel

from prompt_auto_lab.api.models import Session, SessionConfig, SessionStatus, Checkpoint, SessionUpdateRequest
from prompt_auto_lab.api.services import user_service, picasso_service
from prompt_auto_lab.api.services.copilot_queue_service import get_copilot_queue
from prompt_auto_lab.api.websocket.handler import manager as ws_manager
from prompt_auto_lab.api.routes.versions import _load_latest_version
from prompt_auto_lab.core.protocol import (
    GitClient,
    InMemoryEvalBackend,
)
from prompt_auto_lab.core.logging_setup import resolve_log_level
from prompt_auto_lab.core.metrics_config import (
    derive_key_metrics_from_metrics_config,
    normalize_metrics_config,
    resolve_metrics_config,
)
from prompt_auto_lab.core.liquid_entries import resolve_liquid_prompt_entries
from prompt_auto_lab.core.store_dir import get_store_dir
from prompt_auto_lab.orchestrator.orchestrator import Orchestrator, default_store_dir, find_latest_checkpoint

router = APIRouter(prefix="/api/sessions", tags=["sessions"])
logger = logging.getLogger(__name__)


async def _wait_for_copilot_extension_connection(*, timeout_s: float) -> bool:
    """Wait for the VS Code extension to register Copilot models with this server process."""
    queue = get_copilot_queue()
    if queue.is_extension_connected():
        return True

    deadline = time.monotonic() + max(0.0, timeout_s)
    while time.monotonic() < deadline:
        if queue.is_extension_connected():
            return True
        await asyncio.sleep(0.25)

    return queue.is_extension_connected()

_session_user_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "autopromptlab_session_user_id", default=None
)
_session_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "autopromptlab_session_id", default=None
)
_NOISY_LOGGERS = ("httpx", "autogen_core.events")


class _SessionLogFilter(logging.Filter):
    def __init__(self, user_id: str, session_id: str) -> None:
        super().__init__()
        self.user_id = user_id
        self.session_id = session_id

    def filter(self, record: logging.LogRecord) -> bool:
        return _session_user_id.get() == self.user_id and _session_id.get() == self.session_id


def _attach_session_log_handler(
    user_id: str,
    session_id: str,
    log_dir: Path,
) -> tuple[logging.Handler, logging.Handler, Path, Path]:
    log_dir.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
    log_path = log_dir / f"api_{timestamp}.log"
    verbose_path = log_dir / f"api_{timestamp}.verbose.log"
    formatter = logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s")
    log_level = resolve_log_level()
    handler = logging.FileHandler(log_path, encoding="utf-8")
    handler.setFormatter(formatter)
    handler.setLevel(log_level)
    handler.addFilter(_SessionLogFilter(user_id, session_id))
    handler._autopromptlab_session_file = True
    root_logger = logging.getLogger()
    root_logger.setLevel(log_level)
    root_logger.addHandler(handler)
    verbose_handler = logging.FileHandler(verbose_path, encoding="utf-8")
    verbose_handler.setFormatter(formatter)
    verbose_handler.setLevel(log_level)
    verbose_handler.addFilter(_SessionLogFilter(user_id, session_id))
    verbose_handler._autopromptlab_session_verbose_file = True
    for name in _NOISY_LOGGERS:
        noisy_logger = logging.getLogger(name)
        noisy_logger.setLevel(log_level)
        noisy_logger.addHandler(verbose_handler)
    return handler, verbose_handler, log_path, verbose_path


def _detach_session_log_handler(
    handler: logging.Handler | None,
    verbose_handler: logging.Handler | None,
) -> None:
    if handler is not None:
        logging.getLogger().removeHandler(handler)
        handler.close()
    if verbose_handler is None:
        return
    for name in _NOISY_LOGGERS:
        logging.getLogger(name).removeHandler(verbose_handler)
    verbose_handler.close()

class ResumeRequest(PydanticBaseModel):
    checkpoint_path: str | None = None

# Load local_debug from environment (same as main.py)
local_debug_str = os.getenv("LOCAL_DEBUG", "False").lower()
local_debug = local_debug_str in ("true",)
logger.info("Session API local_debug=%s", local_debug)

# In-memory session states and running tasks
_session_states: dict[str, dict[str, Any]] = {}
_running_tasks: dict[str, asyncio.Task] = {}
# User input queues per session - allows WebSocket to pass input to orchestrator
_user_input_queues: dict[str, asyncio.Queue] = {}
# Agent events per session
_session_events: dict[str, list[dict[str, Any]]] = {}
# Event sequence counters per session - ensures correct ordering in UI
_session_event_seq: dict[str, int] = {}
# AIPA eval branch names per session - used for displaying in stop confirmation
_session_aipa_branches: dict[str, str] = {}
# Live eval client per session - used to snapshot active branches for cancellation on stop
_session_github_evals: dict[str, Any] = {}


def get_user_input_queue(session_id: str) -> asyncio.Queue:
    """Get or create a user input queue for a session."""
    if session_id not in _user_input_queues:
        _user_input_queues[session_id] = asyncio.Queue()
    return _user_input_queues[session_id]


async def submit_user_input(session_id: str, message: str) -> None:
    """Submit user input to a session's queue."""
    queue = get_user_input_queue(session_id)
    await queue.put(message)


def _get_next_event_seq(session_id: str) -> int:
    """Get the next sequence number for a session's events."""
    _session_event_seq[session_id] = _session_event_seq.get(session_id, 0) + 1
    return _session_event_seq[session_id]


async def _broadcast_agent_event(
    session_id: str, event_type: str, agent: str, message: str, payload: dict = None
) -> None:
    """Broadcast an agent event with sequence number (for use outside orchestrator context).

    This function is designed to be failure-tolerant - WebSocket failures will not
    interrupt the orchestrator. Events are always persisted to disk first.
    """
    event = {
        "type": event_type,
        "agent": agent,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "seq": _get_next_event_seq(session_id),
        "session_id": session_id,  # Include session_id for UI filtering
        "payload": {"message": message, **(payload or {})},
    }
    # Persist event to file first (critical path)
    _append_session_event(session_id, event)
    # Broadcast to WebSocket (non-critical, fire-and-forget)
    try:
        await ws_manager.broadcast_to_session(
            session_id,
            {
                "type": "agent_event",
                "payload": event,
            },
        )
    except Exception as e:
        # Log but don't fail - UI can poll for events
        logger.debug("WebSocket broadcast failed (non-critical): %s", e)


async def setup_aipa_eval_branch(aipa_git_client: GitClient, base_branch: str, run_id: str) -> tuple[str, bool, str | None]:
    """
    Set up a new branch for aipa-evals based on the specified base branch.

    Args:
        aipa_git_client: GitClient instance for the aipa-evals repository
        base_branch: The base branch to create the new branch from
        run_id: The run ID to include in the branch name

    Returns:
        Tuple of (new_branch_name, success, error_message)
        - new_branch_name: The created branch name or base_branch on failure
        - success: True if branch setup succeeded, False otherwise
        - error_message: Detailed error message if failed, None if successful
    """
    from prompt_auto_lab.core.orchestrator_builder import setup_aipa_eval_branch as shared_setup_aipa_eval_branch

    return await shared_setup_aipa_eval_branch(
        aipa_git_client,
        base_branch=base_branch,
        run_id=run_id,
        logger_=logger,
    )


async def cleanup_aipa_eval_branch(aipa_git_client: GitClient, branch_name: str, base_branch: str) -> bool:
    """
    Clean up the temporary aipa-eval branch after orchestrator run.
    
    Args:
        aipa_git_client: GitClient instance for the aipa-evals repository
        branch_name: The branch name to delete
        base_branch: The base branch to checkout before deleting
    
    Returns:
        True if cleanup was successful, False otherwise
    """
    from prompt_auto_lab.core.orchestrator_builder import cleanup_aipa_eval_branch as shared_cleanup_aipa_eval_branch

    return await shared_cleanup_aipa_eval_branch(
        aipa_git_client,
        branch_name=branch_name,
        base_branch=base_branch,
        logger_=logger,
    )


def _get_session_events_file(session_id: str) -> Path:
    """Get path to session events file."""
    store_dir = get_store_dir() / "events"
    store_dir.mkdir(parents=True, exist_ok=True)
    return store_dir / f"{session_id}.jsonl"


def _load_session_events(session_id: str) -> list[dict[str, Any]]:
    """Load events for a session from file."""
    events_file = _get_session_events_file(session_id)
    events = []
    if events_file.exists():
        try:
            with open(events_file, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        events.append(json.loads(line))
        except Exception:
            pass
    return events


def _append_session_event(session_id: str, event: dict[str, Any]) -> None:
    """Append an event to session's events file."""
    events_file = _get_session_events_file(session_id)
    with open(events_file, "a", encoding="utf-8") as f:
        f.write(json.dumps(event, default=str) + "\n")
    # Also keep in memory (limit to last 500)
    if session_id not in _session_events:
        _session_events[session_id] = []
    _session_events[session_id].append(event)
    if len(_session_events[session_id]) > 500:
        _session_events[session_id] = _session_events[session_id][-500:]


def _get_session_config_file(user_id: str, session_id: str) -> Path:
    """Get path to session config file."""
    session_dir = user_service.get_user_session_dir(user_id, session_id)
    return session_dir / "config.json"


def _save_session_config(user_id: str, session_id: str, config: dict[str, Any]) -> None:
    """Save session config to disk."""
    config_file = _get_session_config_file(user_id, session_id)
    config_file.parent.mkdir(parents=True, exist_ok=True)
    config_file.write_text(json.dumps(config, indent=2, default=str), encoding="utf-8")


def _load_session_config(user_id: str, session_id: str) -> dict[str, Any]:
    """Load session config from disk."""
    config_file = _get_session_config_file(user_id, session_id)
    if not config_file.exists():
        return {}
    try:
        return json.loads(config_file.read_text(encoding="utf-8"))
    except Exception:
        return {}


def _resolve_session_config(user_id: str, session_id: str) -> dict[str, Any]:
    """
    Resolve effective session config by merging persisted config as defaults
    with in-memory runtime state as overrides.

    This avoids losing required static fields (e.g. azure_endpoint) when
    _session_states currently holds only runtime/iteration state.
    """
    state = _session_states.get(session_id)
    state_dict = dict(state) if isinstance(state, dict) else {}
    persisted = _load_session_config(user_id, session_id)

    if persisted:
        merged = dict(persisted)
        merged.update(state_dict)
        _session_states[session_id] = merged
        return merged

    if state_dict:
        _session_states[session_id] = state_dict
        return state_dict

    return {}


def _get_sessions_file(user_id: str) -> Path:
    """Get path to user's sessions storage file."""
    return user_service.get_user_sessions_file(user_id)


def _load_sessions(user_id: str) -> dict[str, Session]:
    """Load sessions for a specific user from their sessions.json."""
    sessions_file = _get_sessions_file(user_id)
    if not sessions_file.exists():
        return {}
    try:
        data = json.loads(sessions_file.read_text(encoding="utf-8"))
        sessions = {}
        for session_data in data:
            session = Session(**session_data)
            sessions[session.id] = session
        return sessions
    except Exception:
        return {}


def _save_sessions(user_id: str, sessions: dict[str, Session]) -> None:
    """Save sessions for a specific user to their sessions.json."""
    sessions_file = _get_sessions_file(user_id)
    sessions_file.parent.mkdir(parents=True, exist_ok=True)
    data = [s.model_dump(mode="json") for s in sessions.values()]
    sessions_file.write_text(json.dumps(data, indent=2, default=str), encoding="utf-8")


def _hydrate_state_from_versions(session_id: str) -> dict[str, Any]:
    """
    Fallback: load latest state snapshot from version.jsonl under the session's store dir.
    Note: The orchestrator writes to 'version.jsonl' (singular).
    """
    versions_file = default_store_dir(session_id) / "version.jsonl"
    if not versions_file.exists():
        return {}
    try:
        lines = [ln for ln in versions_file.read_text(encoding="utf-8").splitlines() if ln.strip()]
        if not lines:
            return {}
        latest = json.loads(lines[-1])
        return latest if isinstance(latest, dict) else {}
    except Exception:
        return {}


def _verify_session_ownership(session_id: str, user_id: str) -> Session:
    """Verify user owns the session. Returns session or raises 404/403."""
    sessions = _load_sessions(user_id)
    if session_id not in sessions:
        raise HTTPException(status_code=404, detail="Session not found")
    session = sessions[session_id]
    if session.user_id != user_id:
        raise HTTPException(status_code=403, detail="Access denied")
    return session


@router.get("")
async def list_sessions(x_user_id: str = Header(None)) -> list[Session]:
    """List all sessions for the current user."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    user = user_service.get_user(x_user_id)
    if user is None:
        raise HTTPException(status_code=401, detail="User not found")

    sessions = _load_sessions(x_user_id)
    return list(sessions.values())


@router.post("")
async def create_session(config: SessionConfig, x_user_id: str = Header(None)) -> Session:
    """Create a new optimization session (instant - repo clone deferred to start)."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    user = user_service.get_user(x_user_id)
    if user is None:
        raise HTTPException(status_code=401, detail="User not found")

    session_id = str(int(time.time())) + "_" + str(uuid.uuid4())
    now = datetime.now(timezone.utc)

    # Create session directory (but NOT picasso subdirectory - clone happens at start)
    session_dir = user_service.get_user_session_dir(x_user_id, session_id)
    session_dir.mkdir(parents=True, exist_ok=True)

    session = Session(
        id=session_id,
        user_id=x_user_id,
        name=f"{config.task_type}_{now.strftime('%Y%m%d_%H%M%S')}",
        status=SessionStatus.IDLE,
        task_type=config.task_type,
        created_at=now,
        updated_at=now,
        current_round=0,
        max_rounds=config.max_rounds,
    )

    # Save to user's sessions
    sessions = _load_sessions(x_user_id)
    sessions[session_id] = session
    _save_sessions(x_user_id, sessions)

    # Store session state in memory (picasso_path computed at start time)
    # Convert legacy liquid_prompt_file_path to liquid_prompt_entries for backward compatibility.
    liquid_entries = resolve_liquid_prompt_entries(
        [e.model_dump() for e in config.liquid_prompt_entries] if config.liquid_prompt_entries else None,
        legacy_liquid_prompt_file_path=config.liquid_prompt_file_path,
    )

    metrics_config = normalize_metrics_config(
        [m.model_dump() for m in config.metrics_config] if config.metrics_config else None
    )
    key_metrics = derive_key_metrics_from_metrics_config(metrics_config)

    session_config = {
        "user_id": x_user_id,
        "user_requirements": config.model_dump(),
        "task_type": config.task_type,
        "goal": config.goal,
        "dataset_name": config.dataset_name,
        "key_metrics": key_metrics,
        "metrics_config": metrics_config,
        "picasso_version": config.picasso_version,
        "conversation_mode": config.conversation_mode,
        "max_dataset_rows": config.max_dataset_rows,
        "eval_project_name": config.eval_project_name,
        "dataset_project_name": config.dataset_project_name,
        "branch": config.branch,
        "max_rounds": config.max_rounds,
        "original_prompt_file_path": config.original_prompt_file_path,
        "liquid_prompt_entries": liquid_entries if liquid_entries else None,
        "termination_rule": config.termination_rule,
        "commit_id": config.commit_id,
        "eval_name": config.eval_name,
        "add_feature_flags": config.add_feature_flags,
        "baseline_feature_flags": config.baseline_feature_flags,
        "model_client_type": config.model_client_type,
        "model_deployment": config.model_deployment,
        "azure_endpoint": config.azure_endpoint,
        "azure_api_version": config.azure_api_version,
        "copilot_model": config.copilot_model,
        "papyrus_model": config.papyrus_model,
        "papyrus_endpoint": config.papyrus_endpoint,
        "papyrus_quota_id": config.papyrus_quota_id,
        "papyrus_scope": config.papyrus_scope,
        "local_model": config.local_model,
        "local_endpoint": config.local_endpoint,
        "selection_strategy": config.selection_strategy,
        "pareto_max_selected": config.pareto_max_selected,
        "baseline_experiment_id": config.baseline_experiment_id,
        "baseline_experiment_link": config.baseline_experiment_link,
        "local_debug": config.local_debug,
        "badcase_deep_analysis": config.badcase_deep_analysis,
        "run_leaderboard_after_optimization": config.run_leaderboard_after_optimization,
        "braintrust_org": config.braintrust_org,
        "environment": config.environment,
        "model": config.model,
        "dataset_sampling_method": config.dataset_sampling_method,
        "additional_cmd_args": config.additional_cmd_args,
        "picasso_env": config.picasso_env,
        "task_id_prefix": config.task_id_prefix,
    }
    _session_states[session_id] = session_config

    # Persist config to disk (survives server restart)
    _save_session_config(x_user_id, session_id, session_config)

    return session


@router.get("/{session_id}")
async def get_session(session_id: str, x_user_id: str = Header(None)) -> Session:
    """Get session details."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")
    return _verify_session_ownership(session_id, x_user_id)


@router.patch("/{session_id}")
async def update_session(
    session_id: str,
    body: SessionUpdateRequest,
    x_user_id: str = Header(None),
) -> Session:
    """Update session properties (e.g., rename)."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    _verify_session_ownership(session_id, x_user_id)

    sessions = _load_sessions(x_user_id)
    session = sessions[session_id]

    if body.name is not None:
        session.name = body.name
        session.updated_at = datetime.now(timezone.utc)

    _save_sessions(x_user_id, sessions)
    return session


@router.get("/{session_id}/checkpoints")
async def list_checkpoints(session_id: str, x_user_id: str = Header(None)) -> list[Checkpoint]:
    """List available checkpoints for a session."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    _verify_session_ownership(session_id, x_user_id)

    # Get checkpoint directory for this session
    session_dir = user_service.get_user_session_dir(x_user_id, session_id)
    checkpoint_dir = session_dir / "checkpoints"

    if not checkpoint_dir.exists():
        return []

    checkpoints = []
    for checkpoint_file in checkpoint_dir.glob("checkpoint_*.json"):
        try:
            # Parse filename: checkpoint_0001_AgentName.json
            name = checkpoint_file.stem
            parts = name.split("_")
            if len(parts) >= 3:
                agent_turn = int(parts[1])
                agent_name = "_".join(parts[2:])  # Handle agent names with underscores

                # Read saved_at from checkpoint file
                data = json.loads(checkpoint_file.read_text(encoding="utf-8"))
                saved_at = data.get("saved_at", "")

                checkpoints.append(Checkpoint(
                    path=str(checkpoint_file.relative_to(session_dir)),
                    agent_turn=agent_turn,
                    agent_name=agent_name,
                    saved_at=saved_at,
                ))
        except (ValueError, json.JSONDecodeError):
            continue

    # Sort by agent_turn descending (latest first)
    checkpoints.sort(key=lambda c: c.agent_turn, reverse=True)
    return checkpoints


@router.post("/{session_id}/resume")
async def resume_session(
    session_id: str,
    body: ResumeRequest | None = None,
    x_user_id: str = Header(None),
) -> Session:
    """Resume a session from a checkpoint."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    session = _verify_session_ownership(session_id, x_user_id)
    if session.status == SessionStatus.RUNNING:
        raise HTTPException(status_code=400, detail="Session already running")

    # Get session directory and checkpoint directory
    session_dir = user_service.get_user_session_dir(x_user_id, session_id)
    checkpoint_dir = session_dir / "checkpoints"

    # Determine checkpoint path
    checkpoint_path: Path | None = None
    if body and body.checkpoint_path:
        checkpoint_path = session_dir / body.checkpoint_path
        if not checkpoint_path.exists():
            raise HTTPException(status_code=400, detail=f"Checkpoint not found: {body.checkpoint_path}")
    else:
        # Find latest checkpoint
        checkpoint_path = find_latest_checkpoint(checkpoint_dir)
        if not checkpoint_path:
            raise HTTPException(status_code=400, detail="No checkpoints found for this session")

    # Update session status
    sessions = _load_sessions(x_user_id)
    sessions[session_id].status = SessionStatus.RUNNING
    sessions[session_id].updated_at = datetime.now(timezone.utc)
    _save_sessions(x_user_id, sessions)

    # Always resolve effective config: persisted defaults + runtime overrides.
    config = _resolve_session_config(x_user_id, session_id)
    if not config:
        raise HTTPException(status_code=400, detail="Session config not found")

    # Start orchestrator in background task with resume_from
    task = asyncio.create_task(_run_orchestrator(session_id, config, resume_from=checkpoint_path))
    _running_tasks[session_id] = task

    return sessions[session_id]


async def _run_orchestrator(
    session_id: str,
    config: dict[str, Any],
    resume_from: str | Path | None = None,
) -> None:
    """Run the orchestrator for a session in the background."""
    user_id = config.get("user_id", "")
    session_dir = user_service.get_user_session_dir(user_id, session_id) if user_id else default_store_dir(session_id)
    session_log_handler: logging.Handler | None = None
    session_log_path: Path | None = None
    session_verbose_handler: logging.Handler | None = None
    session_verbose_path: Path | None = None
    token_user = _session_user_id.set(user_id)
    token_session = _session_id.set(session_id)

    async def broadcast_with_seq(event_type: str, agent: str, message: str, payload: dict = None):
        """Broadcast event with auto-incrementing sequence number."""
        # Use the global session sequence counter for consistent ordering
        await _broadcast_agent_event(session_id, event_type, agent, message, payload)

    try:
        session_log_dir = session_dir / "api_logs"
        (
            session_log_handler,
            session_verbose_handler,
            session_log_path,
            session_verbose_path,
        ) = _attach_session_log_handler(user_id, session_id, session_log_dir)
        logger.info("Session API logging to: %s", session_log_path)
        logger.info("Session API verbose logging to: %s", session_verbose_path)
    except Exception:
        # Don't fail the session if log setup fails.
        session_log_handler = None
        session_log_path = None
        session_verbose_handler = None
        session_verbose_path = None
    logger.info("Session start: user_id=%s session_id=%s", user_id, session_id)
    logger.info("Session config: %s", json.dumps(config, ensure_ascii=False))

    # Progress callback for clone steps
    async def clone_progress(step: int, total: int, message: str) -> None:
        await broadcast_with_seq("agent_message", "System", f"[{step}/{total}] {message}")

    try:
        picasso_path = picasso_service.get_session_picasso_path(user_id, session_id)
        branch = config.get("branch", "main")
        logger.info("Picasso path: %s branch=%s", picasso_path, branch)

        # Step 1: Clone repository (skip if resuming - repo already exists)
        is_resuming = resume_from is not None

        if is_resuming:
            await broadcast_with_seq("agent_message", "System", f"Resuming from checkpoint: {resume_from}")
            # Verify picasso repo exists
            if not picasso_path.exists() or not (picasso_path / ".git").exists():
                raise ValueError(f"Picasso repo not found for resume: {picasso_path}")
        else:
            # Original clone logic - broadcast start message
            await broadcast_with_seq("agent_start", "System", "Cloning repository...")

            # Handle partial/broken clone or skip if already cloned
            try:
                if picasso_path.exists():
                    if (picasso_path / ".git").exists():
                        # Already cloned, skip
                        await broadcast_with_seq("agent_message", "System", "Repository already exists, skipping clone")
                        logger.info("Repository already exists, skipping clone")
                    else:
                        # Partial/broken clone, clean up and re-clone
                        shutil.rmtree(picasso_path)
                        await picasso_service.copy_picasso_repo(
                            picasso_path, branch=branch, progress_callback=clone_progress
                        )
                        await broadcast_with_seq("agent_message", "System", "Repository cloned successfully")
                        logger.info("Repository cloned successfully")
                else:
                    await picasso_service.copy_picasso_repo(
                        picasso_path, branch=branch, progress_callback=clone_progress
                    )
                    await broadcast_with_seq("agent_message", "System", "Repository cloned successfully")
                    logger.info("Repository cloned successfully")
            except Exception as e:
                # Clone failed - broadcast error, revert to IDLE, exit
                await broadcast_with_seq("agent_message", "System", f"Failed to clone repository: {e}")
                logger.error("Clone failed: %s", e)
                # Revert session to IDLE so user can retry
                if user_id:
                    sessions = _load_sessions(user_id)
                    if session_id in sessions:
                        sessions[session_id].status = SessionStatus.IDLE
                        sessions[session_id].updated_at = datetime.now(timezone.utc)
                        _save_sessions(user_id, sessions)
                return

        # Create isolated session branch from user's branch
        if not is_resuming:
            from prompt_auto_lab.core.picasso_workspace import create_session_branch

            await broadcast_with_seq(
                "agent_message", "System",
                f"Creating session branch from {branch}..."
            )
            session_branch = await create_session_branch(picasso_path, session_id, branch)
            config["branch"] = session_branch
            config["session_branch"] = session_branch
            config["user_branch"] = branch
            _session_states[session_id] = config
            if user_id:
                _save_session_config(user_id, session_id, config)
            await broadcast_with_seq(
                "agent_message", "System",
                f"Session branch created and pushed: {session_branch}"
            )

        # Fill baseline commit_id if missing (use latest HEAD of cloned branch)
        if not config.get("commit_id"):
            try:
                result = subprocess.run(
                    ["git", "rev-parse", "HEAD"],
                    cwd=str(picasso_path),
                    capture_output=True,
                    text=True,
                    timeout=10,
                )
                if result.returncode == 0:
                    head_commit = result.stdout.strip()
                    if head_commit:
                        config["commit_id"] = head_commit
                        _session_states[session_id] = config
                        if user_id:
                            _save_session_config(user_id, session_id, config)
                        await broadcast_with_seq(
                            "agent_message",
                            "System",
                            f"No commit_id provided. Using latest commit on branch {branch}: {head_commit[:7]}",
                        )
                else:
                    logger.warning("Failed to resolve HEAD commit: %s", result.stderr)
            except Exception as e:
                logger.warning("Failed to set commit_id from HEAD: %s", e)

        # Step 2: Continue with orchestrator setup
        await broadcast_with_seq("agent_message", "System", "Initializing orchestrator...")

        # Initialize model clients using shared builder
        from prompt_auto_lab.core.orchestrator_builder import (
            ModelClientBuilder,
            ModelClientConfig,
            ProtocolClientBuilder,
            EvalClientConfig,
            BraintrustConfig,
        )
        
        model_client_type = config.get("model_client_type", "Azure")
        await broadcast_with_seq("agent_message", "System", f"Model client type: {model_client_type}")

        # Log selection strategy configuration
        selection_strategy = config.get("selection_strategy", "greedy")
        pareto_max_selected = config.get("pareto_max_selected", 2)
        await broadcast_with_seq("agent_message", "System", f"Selection strategy: {selection_strategy}" + (f" (max_selected={pareto_max_selected})" if selection_strategy == "gepa" else ""))

        model_config = ModelClientConfig(
            client_type=model_client_type,
            model_deployment=config.get("model_deployment"),
            azure_endpoint=config.get("azure_endpoint"),
            azure_api_version=config.get("azure_api_version"),
            papyrus_endpoint=config.get("papyrus_endpoint"),
            papyrus_model=config.get("papyrus_model"),
            papyrus_quota_id=config.get("papyrus_quota_id"),
            papyrus_scope=config.get("papyrus_scope"),
            copilot_model=config.get("copilot_model", ""),
            local_model=config.get("local_model", ""),
            local_endpoint=config.get("local_endpoint", "http://localhost:23333/api/openai/v1"),
        )

        # Send initialization message based on client type
        if model_client_type == "Copilot":
            await broadcast_with_seq("agent_message", "System", f"Initializing Copilot proxy client with model: '{model_config.copilot_model}'...")
            # Special handling for Copilot - need to wait for extension
            try:
                queue = get_copilot_queue()
                if not queue.is_extension_connected():
                    timeout_s = float(os.getenv("AUTOPROMPTLAB_COPILOT_CONNECT_TIMEOUT_S") or 20.0)
                    await broadcast_with_seq(
                        "agent_message",
                        "System",
                        f"Waiting for VS Code Copilot proxy to connect (up to {timeout_s:.0f}s)...",
                    )
                    ok = await _wait_for_copilot_extension_connection(timeout_s=timeout_s)
                    if not ok:
                        raise RuntimeError(
                            "Copilot provider requires VS Code extension connection. "
                            "Please open VS Code with the AutoPromptLab extension active, then start the Copilot proxy in the sidebar."
                        )
            except RuntimeError as e:
                await broadcast_with_seq("agent_message", "System", f"Copilot initialization failed: {e}")
                raise HTTPException(status_code=400, detail=str(e))
        elif model_client_type == "Papyrus":
            await broadcast_with_seq("agent_message", "System", "Initializing Papyrus credentials...")
        elif model_client_type == "Local":
            await broadcast_with_seq("agent_message", "System", f"Initializing Local client with model: '{model_config.local_model}' at {model_config.local_endpoint}...")
        else:
            # Azure
            await broadcast_with_seq("agent_message", "System", "Initializing Azure credentials...")

        try:
            built_clients = ModelClientBuilder.build(model_config)
            if hasattr(built_clients, "main") and hasattr(built_clients, "think"):
                model_client = built_clients.main
                model_think_client = built_clients.think
            else:
                # Backward compatibility for tuple return shape.
                model_client, model_think_client = built_clients
            # Show model-specific success message
            if model_client_type == "Copilot":
                await broadcast_with_seq("agent_message", "System", f"Copilot client initialized. Using model: '{model_client._model}'")
            elif model_client_type == "Local":
                await broadcast_with_seq("agent_message", "System", f"Local client initialized. Using model: '{model_config.local_model}'")
            else:
                await broadcast_with_seq("agent_message", "System", f"{model_client_type} client initialized successfully")
        except Exception as e:
            await broadcast_with_seq("agent_message", "System", f"Model client initialization failed: {e}")
            raise HTTPException(status_code=400, detail=str(e))

        await broadcast_with_seq("agent_message", "System", "Initializing clients...")

        # Initialize clients
        backend = InMemoryEvalBackend()

        # Get session-specific picasso path and output directory
        user_id = config.get("user_id", "")
        output_dir = user_service.get_user_session_dir(user_id, session_id) if user_id else default_store_dir(session_id)

        aipa_eval_branch = config.get("session_branch") or config.get("branch", "main")
        aipa_eval_base_branch = config.get("user_branch") or config.get("branch", "main")
        aipa_branch_created = False
        _session_aipa_branches[session_id] = aipa_eval_branch

        # Use local_debug from session config, fallback to environment
        session_local_debug = config.get("local_debug")
        if session_local_debug is None:
            local_debug_str = os.getenv("LOCAL_DEBUG", "False").lower()
            session_local_debug = local_debug_str in ("true",)

        raw_max_dataset_rows = config.get("max_dataset_rows")
        max_dataset_rows = None
        if isinstance(raw_max_dataset_rows, int):
            max_dataset_rows = raw_max_dataset_rows
        elif isinstance(raw_max_dataset_rows, str):
            try:
                max_dataset_rows = int(raw_max_dataset_rows.strip())
            except ValueError:
                max_dataset_rows = None
        conversation_mode = config.get("conversation_mode") or os.getenv("CONVERSATION_MODE", "DefaultMode")

        # Compute default github_cli_path: prompt_auto_lab/tools/trigger_picasso_eval.ps1
        default_github_cli_path = str(Path(__file__).parent.parent.parent / "tools" / "trigger_picasso_eval.ps1")
        github_eval = ProtocolClientBuilder.build_eval_client(
            EvalClientConfig(
                github_cli_path=os.getenv("GITHUB_CLI_PATH", default_github_cli_path),
                eval_name=config.get("eval_name"),
                dataset_name=config.get("dataset_name", ""),
                output_dir=str(output_dir),
                local_debug=session_local_debug,
                add_feature_flags=config.get("add_feature_flags"),
                baseline_feature_flags=config.get("baseline_feature_flags"),
                aipa_eval_branch=aipa_eval_branch,
                conversation_mode=conversation_mode,
                max_dataset_rows=max_dataset_rows,
                dataset_sampling_method=config.get("dataset_sampling_method"),
                additional_cmd_args=config.get("additional_cmd_args"),
                picasso_version=config.get("picasso_version"),
                eval_project_name=config.get("eval_project_name"),
                dataset_project_name=config.get("dataset_project_name"),
                braintrust_org=config.get("braintrust_org", "non-pme"),
                environment=config.get("environment"),
                model=config.get("model"),
                picasso_env=config.get("picasso_env"),
                task_id_prefix=config.get("task_id_prefix"),
            ),
            backend,
        )
        _session_github_evals[session_id] = github_eval

        # Use picasso_path computed at clone step (not from config)
        if not picasso_path.exists():
            raise ValueError(f"Session picasso path not found: {picasso_path}")

        # Resolve metrics_config/key_metrics with shared normalization + dataset defaults.
        dataset_name = config.get("dataset_name", "")
        metrics_config, key_metrics = resolve_metrics_config(
            dataset_name=dataset_name,
            metrics_config=config.get("metrics_config"),
        )

        braintrust_org = config.get("braintrust_org", "non-pme")
        braintrust_api_key = os.getenv("BRAINTRUST_AZURE_API_KEY", "") if braintrust_org == "bt-azure" else os.getenv("BRAINTRUST_API_KEY", "")

        # Determine model endpoint and name based on client type for BTQL agent subprocess
        btql_model_endpoint = None
        btql_model_name = None
        if model_client_type == "Papyrus":
            btql_model_endpoint = config.get("papyrus_endpoint") or os.getenv("PAPYRUS_ENDPOINT")
            btql_model_name = config.get("papyrus_model") or os.getenv("PAPYRUS_MODEL")
        elif model_client_type == "Copilot":
            btql_model_name = config.get("copilot_model")
        elif model_client_type == "Local":
            btql_model_endpoint = config.get("local_endpoint")
            btql_model_name = config.get("local_model")
        # Azure uses environment variables directly in btql-agent

        braintrust = ProtocolClientBuilder.build_braintrust_client(
            BraintrustConfig(
                api_key=braintrust_api_key,
                metrics_config=metrics_config,  # Pass metrics configuration
                output_dir=os.path.join(str(output_dir), "braintrust_outputs"),
                local_debug=session_local_debug,
                eval_name=config.get("eval_name"),
                braintrust_org=braintrust_org,
                environment=config.get("environment"),
                model_client_type=model_client_type.lower(),
                model_endpoint=btql_model_endpoint,
                model_name=btql_model_name,
                log_callback=broadcast_with_seq,
                badcase_deep_analysis=bool(config.get("badcase_deep_analysis", False)),
            ),
            backend,
            model_client,
        )
        picasso = ProtocolClientBuilder.build_picasso_client(picasso_path)
        git_agent = ProtocolClientBuilder.build_git_client(picasso_path, output_dir)
        liquid_executor, python_executor = ProtocolClientBuilder.build_code_executors()

        # Get user input queue for this session
        user_input_queue = get_user_input_queue(session_id)

        # Disable structured output for Local and Copilot models
        use_structured_output = model_client_type not in ("Local", "Copilot")

        orchestrator = Orchestrator(
            model_client=model_client,
            model_think_client=model_think_client,
            aipa=github_eval,
            braintrust=braintrust,
            picasso=picasso,
            git_agent=git_agent,
            liquid_executor=liquid_executor,
            python_executor=python_executor,
            store_dir=output_dir,
            use_structured_output=use_structured_output,
        )

        # Set the user input queue
        orchestrator._user_input_queue = user_input_queue
        orchestrator._user_proxy_mode = "web"

        # Prepare user input using shared builder
        from prompt_auto_lab.core.orchestrator_builder import UserInputBuilder
        user_input = UserInputBuilder.build(
            task_type=config.get("task_type", "picasso"),
            goal=config.get("goal", "Optimize a base prompt for higher evaluation score."),
            dataset_name=config.get("dataset_name", ""),
            key_metrics=key_metrics,
            eval_name=config.get("eval_name"),
            eval_project_name=config.get("eval_project_name"),
            dataset_project_name=config.get("dataset_project_name"),
            braintrust_org=braintrust_org,
            picasso_version=config.get("picasso_version"),
            conversation_mode=config.get("conversation_mode"),
            max_dataset_rows=config.get("max_dataset_rows"),
            branch=config.get("branch", "main"),
            commit_id=config.get("commit_id"),
            liquid_prompt_entries=config.get("liquid_prompt_entries"),
            original_prompt_file_path=config.get("original_prompt_file_path"),
            termination_rule=config.get("termination_rule", {"type": "no_progress", "rounds": 3}),
            baseline_experiment_id=config.get("baseline_experiment_id"),
            baseline_experiment_link=config.get("baseline_experiment_link"),
            metrics_config=metrics_config,
            environment=config.get("environment"),
            model=config.get("model"),
            run_leaderboard_after_optimization=config.get("run_leaderboard_after_optimization"),
            max_rounds=config.get("max_rounds"),
            selection_strategy=config.get("selection_strategy", "greedy"),
            pareto_max_selected=config.get("pareto_max_selected", 2),
        )

        # Store callback in orchestrator for use during execution
        # Use broadcast_with_seq which includes sequence numbers for correct ordering
        orchestrator._broadcast_event = broadcast_with_seq
        orchestrator._session_id = session_id

        # Get checkpoint directory
        checkpoint_dir = output_dir / "checkpoints"

        # Run orchestrator
        try:
            logger.info("Starting orchestrator run")
            await orchestrator.run(
                user_input=user_input,
                max_rounds=config.get("max_rounds", 5),
                use_console=True,
                checkpoint_dir=checkpoint_dir,
                checkpoint_each_turn=True,
                resume_from=resume_from,
            )
            logger.info("Orchestrator run completed")

            # After orchestrator completes, trigger leaderboard eval with best branch if enabled
            try:
                from prompt_auto_lab.core.orchestrator_builder import submit_leaderboard_job

                user_id = config.get("user_id", "")
                latest_version = _load_latest_version(user_id, session_id) if user_id else None
                iteration = (latest_version or {}).get("iteration") if isinstance(latest_version, dict) else None

                async def _notify(message: str) -> None:
                    await broadcast_with_seq("agent_message", "System", message)

                await submit_leaderboard_job(
                    enabled=bool(config.get("run_leaderboard_after_optimization", False)),
                    github_eval_client=github_eval,
                    aipa_eval_branch=aipa_eval_branch,
                    iteration=iteration if isinstance(iteration, dict) else None,
                    logger_=logger,
                    notify=_notify,
                )
            except Exception as e:
                logger.warning("Failed to trigger leaderboard eval: %s", e)
                # Don't fail the session if leaderboard eval fails
        except asyncio.CancelledError:
            # Session was stopped by user
            await broadcast_with_seq("agent_complete", "Orchestrator", "Orchestrator cancelled by user")
            logger.info("Orchestrator cancelled by user")
            raise  # Re-raise so the task completes
        finally:
            from prompt_auto_lab.core.orchestrator_builder import finalize_aipa_task

            async def _async_cleanup():
                async def _notify(message: str) -> None:
                    await broadcast_with_seq("agent_message", "System", message)

                await finalize_aipa_task(
                    github_eval_client=github_eval,
                    aipa_eval_branch=aipa_eval_branch,
                    aipa_eval_base_branch=aipa_eval_base_branch,
                    aipa_branch_created=aipa_branch_created,
                    logger_=logger,
                    notify=_notify,
                    skip_cleanup_on_main=True,
                )
            
            try:
                await asyncio.shield(_async_cleanup())
            except asyncio.CancelledError:
                # Cleanup was interrupted, but we've shielded it so it should complete
                pass
            logger.info("Async cleanup completed")

        # Update session status on completion
        user_id = config.get("user_id", "")
        if user_id:
            sessions = _load_sessions(user_id)
            if session_id in sessions:
                sessions[session_id].status = SessionStatus.COMPLETED
                sessions[session_id].updated_at = datetime.now(timezone.utc)
                _save_sessions(user_id, sessions)

        await broadcast_with_seq("agent_complete", "Orchestrator", "Optimization loop completed")
        logger.info("Session completed: user_id=%s session_id=%s", user_id, session_id)

        await model_client.close()

    except Exception as e:
        logger.exception("Orchestrator error")

        # Update session status on error
        user_id = config.get("user_id", "")
        if user_id:
            sessions = _load_sessions(user_id)
            if session_id in sessions:
                sessions[session_id].status = SessionStatus.ERROR
                sessions[session_id].updated_at = datetime.now(timezone.utc)
                _save_sessions(user_id, sessions)

        await broadcast_with_seq("agent_message", "System", f"Error: {str(e)}")
        await ws_manager.broadcast_to_session(
            session_id,
            {
                "type": "error",
                "payload": {"message": f"Orchestrator error: {str(e)}"},
            },
        )
    finally:
        _session_github_evals.pop(session_id, None)
        _detach_session_log_handler(session_log_handler, session_verbose_handler)
        _session_user_id.reset(token_user)
        _session_id.reset(token_session)


@router.post("/{session_id}/start")
async def start_session(session_id: str, x_user_id: str = Header(None)) -> Session:
    """Start a session (trigger orchestrator run)."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    session = _verify_session_ownership(session_id, x_user_id)
    if session.status == SessionStatus.RUNNING:
        raise HTTPException(status_code=400, detail="Session already running")
    logger.info("Start session request: user_id=%s session_id=%s", x_user_id, session_id)

    # Update session status
    sessions = _load_sessions(x_user_id)
    sessions[session_id].status = SessionStatus.RUNNING
    sessions[session_id].updated_at = datetime.now(timezone.utc)
    _save_sessions(x_user_id, sessions)

    # Always resolve effective config: persisted defaults + runtime overrides.
    config = _resolve_session_config(x_user_id, session_id)
    if not config:
        raise HTTPException(status_code=400, detail="Session config not found")

    # Start orchestrator in background task
    task = asyncio.create_task(_run_orchestrator(session_id, config))
    _running_tasks[session_id] = task

    return sessions[session_id]


@router.get("/{session_id}/aipa_branch")
async def get_session_aipa_branch(session_id: str, x_user_id: str = Header(None)) -> dict:
    """Get the aipa-eval branch name and currently active eval branches for a session."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    _verify_session_ownership(session_id, x_user_id)
    branch_name = _session_aipa_branches.get(session_id)
    eval_client = _session_github_evals.get(session_id)
    active: list[str] = sorted(eval_client._active_branches) if eval_client and eval_client._active_branches else []
    return {"aipa_branch": branch_name, "active_branches": active}


@router.post("/{session_id}/stop")
async def stop_session(session_id: str, x_user_id: str = Header(None)) -> Session:
    """Stop a running session."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    _verify_session_ownership(session_id, x_user_id)
    logger.info("Stop session request: user_id=%s session_id=%s", x_user_id, session_id)

    # Cancel any in-progress GitHub Actions workflows on all active eval branches
    # Snapshot active branches BEFORE cancelling the asyncio task so we capture the
    # branch that submit_eval is currently running on (its finally block will clear
    # _active_branches once CancelledError propagates through).
    eval_client = _session_github_evals.get(session_id)
    active_branches: set[str] = set(eval_client._active_branches) if eval_client else set()
    # Fallback to session branch if nothing is actively tracked yet
    if not active_branches:
        session_branch = _session_aipa_branches.get(session_id)
        if session_branch:
            active_branches = {session_branch}

    # Kill any running eval subprocess immediately so it cannot retry cancelled workflows
    if eval_client is not None and hasattr(eval_client, "kill_subprocess"):
        try:
            eval_client.kill_subprocess()
        except Exception:
            logger.debug("Failed to kill eval subprocess in stop_session", exc_info=True)

    # Cancel the running task if it exists
    task = _running_tasks.get(session_id)
    if task and not task.done():
        task.cancel()
        try:
            await asyncio.wait_for(task, timeout=5.0)
        except (asyncio.CancelledError, asyncio.TimeoutError):
            pass

    # Cancel any in-progress GitHub Actions workflows on all active eval branches
    if active_branches:
        cancel_eval_client = eval_client
        if cancel_eval_client is None:
            # Fallback: build a minimal eval client just for cancellation
            try:
                from prompt_auto_lab.core.protocol import GitHubActionEvalClient, InMemoryEvalBackend
                default_github_cli_path = str(Path(__file__).parent.parent.parent / "tools" / "trigger_picasso_eval.ps1")
                cancel_eval_client = GitHubActionEvalClient(
                    backend=InMemoryEvalBackend(),
                    github_cli_path=os.getenv("GITHUB_CLI_PATH", default_github_cli_path),
                )
            except Exception:
                cancel_eval_client = None
        if cancel_eval_client is not None:
            for branch in active_branches:
                try:
                    cancel_result = await cancel_eval_client.cancel_all_runs_on_branch(branch)
                    cancelled_count = cancel_result.get("cancelled_count", 0)
                    if cancelled_count > 0:
                        logger.info("Cancelled %d workflow run(s) on branch %s", cancelled_count, branch)
                        await _broadcast_agent_event(session_id, "agent_message", "System", f"Cancelled {cancelled_count} running eval job(s) on branch: {branch}")
                    if cancel_result.get("errors"):
                        logger.warning("Errors cancelling workflows on %s: %s", branch, cancel_result["errors"])
                except Exception:
                    logger.warning("Failed to cancel GitHub Actions workflows on branch %s", branch, exc_info=True)

    # Clean up task and queue
    _running_tasks.pop(session_id, None)
    _user_input_queues.pop(session_id, None)
    _session_aipa_branches.pop(session_id, None)

    # Broadcast stop message
    await _broadcast_agent_event(session_id, "agent_message", "System", "Session stopped by user")

    sessions = _load_sessions(x_user_id)
    sessions[session_id].status = SessionStatus.COMPLETED
    sessions[session_id].updated_at = datetime.now(timezone.utc)
    _save_sessions(x_user_id, sessions)

    return sessions[session_id]


@router.post("/{session_id}/input")
async def send_user_input(session_id: str, body: dict[str, Any], x_user_id: str = Header(None)) -> dict[str, str]:
    """Send user input to a running session."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    session = _verify_session_ownership(session_id, x_user_id)
    if session.status != SessionStatus.RUNNING:
        raise HTTPException(status_code=400, detail="Session is not running")

    # Handle confirmation responses (from VS Code extension dialog)
    if "confirmed" in body:
        # Forward the confirmation as JSON string
        message = json.dumps({"confirmed": body["confirmed"]})
    else:
        # Regular user message
        message = body.get("message", "")
        if not message:
            raise HTTPException(status_code=400, detail="Message or confirmed field is required")

    # Submit to the user input queue
    await submit_user_input(session_id, message)

    # Only broadcast for regular user messages, not confirmation responses
    if "confirmed" not in body:
        await _broadcast_agent_event(session_id, "agent_message", "UserProxy", f"User: {message}")

    return {"status": "ok", "message": "Input submitted"}


@router.get("/{session_id}/state")
async def get_session_state(session_id: str, x_user_id: str = Header(None)) -> dict[str, Any]:
    """Get current SharedState for a session."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    _verify_session_ownership(session_id, x_user_id)

    state = _session_states.get(session_id)
    if not state:
        state = _hydrate_state_from_versions(session_id)
        _session_states[session_id] = state

    # Always merge persisted config as defaults so duplication/UI can reliably
    # access full session config even when state is rehydrated from version.jsonl.
    state = _resolve_session_config(x_user_id, session_id)

    # Update current_round from iteration count if available
    try:
        iteration = state.get("iteration") if isinstance(state, dict) else None
        if isinstance(iteration, dict):
            current_round = len(iteration)
            sessions = _load_sessions(x_user_id)
            if session_id in sessions:
                sessions[session_id].current_round = current_round
                _save_sessions(x_user_id, sessions)
    except Exception:
        pass

    return state or {}


@router.get("/{session_id}/events")
async def get_session_events(
    session_id: str,
    limit: int = 500,
    offset: int = 0,
    since_seq: int = 0,
    x_user_id: str = Header(None),
) -> dict[str, Any]:
    """Get agent events for a session.

    Args:
        session_id: The session ID
        limit: Maximum number of events to return (default 500)
        offset: Number of events to skip (for pagination)
        since_seq: Only return events with seq > since_seq (for polling)
        x_user_id: User ID header for authentication

    Returns:
        dict with events, total count, and pagination info
    """
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    _verify_session_ownership(session_id, x_user_id)

    # Load from file if not in memory
    if session_id not in _session_events:
        _session_events[session_id] = _load_session_events(session_id)

    events = _session_events.get(session_id, [])

    # Filter by since_seq if provided (for polling)
    if since_seq > 0:
        events = [e for e in events if e.get("seq", 0) > since_seq]

    total = len(events)
    paginated = events[offset:offset + limit]

    # Include max_seq for polling - client uses this for next since_seq
    max_seq = max((e.get("seq", 0) for e in paginated), default=since_seq)

    return {
        "events": paginated,
        "total": total,
        "limit": limit,
        "offset": offset,
        "since_seq": since_seq,
        "max_seq": max_seq,
    }


@router.patch("/{session_id}/state")
async def update_session_state(session_id: str, patch: dict[str, Any], x_user_id: str = Header(None)) -> dict[str, Any]:
    """Update SharedState fields."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    _verify_session_ownership(session_id, x_user_id)

    if session_id not in _session_states:
        _session_states[session_id] = _load_session_config(x_user_id, session_id) or {}

    _session_states[session_id].update(patch)
    # Persist patched state so manual edits survive server restart.
    _save_session_config(x_user_id, session_id, _session_states[session_id])
    return _session_states[session_id]


def _get_branches_to_delete(session_id: str, x_user_id: str) -> set[str]:
    """
    Get list of GA branches that will be deleted for a session.
    Returns a set of branch names (excluding user's original branch).
    """
    # Get session config to find the user's original branch (should NOT be deleted)
    session_config = _session_states.get(session_id)
    if not session_config:
        session_config = _load_session_config(x_user_id, session_id)
    user_original_branch = (session_config.get("user_branch") or session_config.get("branch")) if session_config else None

    from prompt_auto_lab.core.orchestrator_builder import collect_candidate_branches_from_task_dir

    session_dir = user_service.get_user_session_dir(x_user_id, session_id)
    picasso_path = picasso_service.get_session_picasso_path(x_user_id, session_id)
    return collect_candidate_branches_from_task_dir(
        session_dir,
        original_branch=user_original_branch,
        repo_path=picasso_path,
    )


@router.get("/{session_id}/branches-to-delete")
async def preview_delete_branches(session_id: str, x_user_id: str = Header(None)) -> dict[str, Any]:
    """Preview which branches will be deleted when deleting this session."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    _verify_session_ownership(session_id, x_user_id)

    branches = _get_branches_to_delete(session_id, x_user_id)

    # Include session branch in preview
    session_config = _session_states.get(session_id) or _load_session_config(x_user_id, session_id)
    session_branch = session_config.get("session_branch") if isinstance(session_config, dict) else None
    if session_branch:
        branches.add(session_branch)

    return {
        "session_id": session_id,
        "branches_to_delete": sorted(list(branches)),
        "count": len(branches),
    }


@router.delete("/{session_id}")
async def delete_session(session_id: str, x_user_id: str = Header(None)) -> dict[str, Any]:
    """Delete a session and all its data, including GA branches (but not the user's original branch)."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    _verify_session_ownership(session_id, x_user_id)

    # Get branches to delete using the helper function
    branches_to_delete = _get_branches_to_delete(session_id, x_user_id)
    session_config = _session_states.get(session_id) or _load_session_config(x_user_id, session_id)
    # Include session branch in cleanup
    session_branch = session_config.get("session_branch") if isinstance(session_config, dict) else None
    if session_branch:
        branches_to_delete.add(session_branch)
    user_original_branch = (session_config.get("user_branch") or session_config.get("branch")) if isinstance(session_config, dict) else None
    session_dir = user_service.get_user_session_dir(x_user_id, session_id)

    # Delete GA branches and session directory via shared helper
    picasso_path = picasso_service.get_session_picasso_path(x_user_id, session_id)
    git_client = None
    if branches_to_delete and picasso_path.exists():
        try:
            from prompt_auto_lab.core.orchestrator_builder import ProtocolClientBuilder

            git_client = ProtocolClientBuilder.build_git_client(picasso_path)
        except Exception as e:
            logger.warning("Failed to initialize GitClient for branch cleanup: %s", e)

    from prompt_auto_lab.core.orchestrator_builder import delete_task_artifacts

    cleanup_result = await delete_task_artifacts(
        task_dir=session_dir,
        original_branch=user_original_branch,
        session_repo_path=picasso_path,
        delete_branches=git_client is not None,
        git_client=git_client,
        branches_to_delete=branches_to_delete,
        logger_=logger,
    )
    session_repo_deleted = cleanup_result.get("session_repo_deleted", False)
    deleted_branches = cleanup_result.get("deleted_branches", [])
    failed_branches = cleanup_result.get("failed_branches", [])
    skipped_branches = cleanup_result.get("skipped_branches", [])
    # Remove from sessions.json
    sessions = _load_sessions(x_user_id)
    sessions.pop(session_id, None)
    _save_sessions(x_user_id, sessions)

    # Clean up in-memory state
    _session_states.pop(session_id, None)

    return {
        "status": "deleted",
        "session_id": session_id,
        "session_repo_deleted": session_repo_deleted,
        "deleted_branches": deleted_branches,
        "failed_branches": failed_branches,
        "skipped_branches": skipped_branches,
    }


@router.post("/{session_id}/complete")
async def complete_session(session_id: str, x_user_id: str = Header(None)) -> Session:
    """Mark session as complete and delete Picasso repo (keep results)."""
    if not x_user_id:
        raise HTTPException(status_code=401, detail="X-User-Id header required")

    _verify_session_ownership(session_id, x_user_id)

    # Update status
    sessions = _load_sessions(x_user_id)
    sessions[session_id].status = SessionStatus.COMPLETED
    sessions[session_id].updated_at = datetime.now(timezone.utc)
    _save_sessions(x_user_id, sessions)

    # Delete only the picasso repo, keep other data
    picasso_path = picasso_service.get_session_picasso_path(x_user_id, session_id)
    picasso_service.delete_picasso_repo(picasso_path)

    return sessions[session_id]
