from __future__ import annotations
import os
import uuid
import logging
import asyncio
import atexit
import multiprocessing
from typing import Any, Dict, List, Optional, Callable, Awaitable
from multiprocessing.connection import Connection
from datetime import datetime, timezone

from mcp_stata.models import SessionInfo, CommandResponse

logger = logging.getLogger("mcp_stata.sessions")

# Use 'spawn' for process creation to ensure thread-safety and Stata/Rust compatibility.
# Re-exposed at module level so tests can patch these references.
_ctx = multiprocessing.get_context("spawn")
Process = _ctx.Process
Pipe = _ctx.Pipe

class StataSession:
    def __init__(self, session_id: str):
        self.id = session_id
        self.status = "starting"
        self.created_at = datetime.now(timezone.utc).isoformat()
        self.pid: Optional[int] = None
        
        self._parent_conn, self._child_conn = Pipe()
        self._process = Process(target=self._run_worker, args=(self._child_conn,))
        self._process.daemon = True
        self._process.start()
        
        self._pending_requests: Dict[str, asyncio.Future] = {}
        self._log_listeners: Dict[str, List[Callable[[str], Awaitable[None]]]] = {}
        self._progress_listeners: Dict[str, List[Callable[[float, Optional[float], Optional[str]], Awaitable[None]]]] = {}
        
        self._listener_running = True
        self._listener_task = asyncio.create_task(self._listen_to_worker())

    def _run_worker(self, conn: Connection):
        from mcp_stata.worker import main
        main(conn)

    async def _listen_to_worker(self):
        loop = asyncio.get_running_loop()
        try:
            while self._listener_running:
                # Use poll with timeout to allow checking self._listener_running and asyncio cancellation
                if await loop.run_in_executor(None, self._parent_conn.poll, 0.2):
                    try:
                        msg = await loop.run_in_executor(None, self._parent_conn.recv)
                        await self._handle_worker_msg(msg)
                    except (EOFError, ConnectionResetError, BrokenPipeError):
                        logger.info(f"Session {self.id} worker connection closed.")
                        break
                else:
                    # Give the event loop a chance to process other tasks and check cancellation
                    await asyncio.sleep(0.01)
        except asyncio.CancelledError:
            self._listener_running = False
            raise
        except Exception as e:
            logger.error(f"Error in session {self.id} listener: {e}")
            self.status = "error"
        finally:
            self._listener_running = False
            if self.status != "error":
                self.status = "stopped"

    async def _handle_worker_msg(self, msg: Dict[str, Any]):
        event = msg.get("event")
        msg_id = msg.get("id")
        
        if event == "ready":
            self.pid = msg.get("pid")
            self.status = "running"
            logger.info(f"Session {self.id} ready (PID: {self.pid})")
        
        elif event == "log":
            if msg_id in self._log_listeners:
                for cb in self._log_listeners[msg_id]:
                    await cb(msg.get("text"))
        
        elif event == "progress":
            if msg_id in self._progress_listeners:
                for cb in self._progress_listeners[msg_id]:
                    await cb(msg.get("progress"), msg.get("total"), msg.get("message"))
        
        elif event == "result":
            if msg_id in self._pending_requests:
                if not self._pending_requests[msg_id].done():
                    self._pending_requests[msg_id].set_result(msg.get("result"))
                self._cleanup_listeners(msg_id)
        
        elif event == "error":
            if msg_id in self._pending_requests:
                if not self._pending_requests[msg_id].done():
                    self._pending_requests[msg_id].set_exception(RuntimeError(msg.get("message")))
                self._cleanup_listeners(msg_id)
            else:
                logger.error(f"Global worker error in session {self.id}: {msg.get('message')}")
                # Don't update status if already stopped or error
                if self.status not in ("stopped", "error"):
                    self.status = "error"

    def _cleanup_listeners(self, msg_id: str):
        self._log_listeners.pop(msg_id, None)
        self._progress_listeners.pop(msg_id, None)
        self._pending_requests.pop(msg_id, None)

    async def _ensure_listener(self):
        current_loop = asyncio.get_running_loop()
        if self._listener_task is None or self._listener_task.done() or (hasattr(self._listener_task, "get_loop") and self._listener_task.get_loop() != current_loop):
            if self._listener_task and not self._listener_task.done():
                self._listener_task.cancel()
            self._listener_running = True
            self._listener_task = current_loop.create_task(self._listen_to_worker())

    async def call(self, method: str, args: Dict[str, Any], 
                   notify_log: Optional[Callable[[str], Awaitable[None]]] = None,
                   notify_progress: Optional[Callable[[float, Optional[float], Optional[str]], Awaitable[None]]] = None) -> Any:
        
        await self._ensure_listener()
        msg_id = uuid.uuid4().hex
        future = asyncio.get_running_loop().create_future()
        self._pending_requests[msg_id] = future
        
        if notify_log:
            self._log_listeners.setdefault(msg_id, []).append(notify_log)
        if notify_progress:
            self._progress_listeners.setdefault(msg_id, []).append(notify_progress)
            
        try:
            self._parent_conn.send({
                "type": method,
                "id": msg_id,
                "args": args
            })
            return await future
        except asyncio.CancelledError:
            # If the session call is cancelled (e.g., from the server or UI),
            # send an out-of-band 'break' message to the worker to interrupt Stata.
            logger.info(f"Cancellation requested for command {method}:{msg_id} in session {self.id}")
            try:
                self._parent_conn.send({"type": "break"})
            except Exception as e:
                logger.warning(f"Failed to send break command to worker for session {self.id}: {e}")
            
            # Wait briefly for the worker to return the result of the interrupted command.
            # We use shield so that we stay in this call until we see the worker acknowledge or we timeout.
            # This prevents the next command from being sent while the worker is still cleaning up.
            try:
                # Give Stata a few seconds to acknowledge the break.
                await asyncio.wait_for(asyncio.shield(future), timeout=3.0)
                logger.info(f"Session {self.id} acknowledged break for {msg_id}")
            except (asyncio.TimeoutError, Exception) as e:
                logger.warning(f"Session {self.id} did not acknowledge break within timeout: {e}")
            
            # Re-raise cancellation
            raise
        except (AttributeError, BrokenPipeError, ConnectionResetError) as e:
             self._cleanup_listeners(msg_id)
             raise RuntimeError(f"Failed to send command to worker: {e}")

    async def send_break(self):
        """Send an out-of-band break signal to the worker."""
        try:
            self._parent_conn.send({"type": "break"})
            logger.info(f"Break signal sent to session {self.id}")
        except Exception as e:
            logger.warning(f"Failed to send break command to session {self.id}: {e}")
            raise RuntimeError(f"Failed to send break signal: {e}")

    async def stop(self, timeout: float = 5.0):
        self._listener_running = False
        if self.status != "stopped":
            try:
                self._parent_conn.send({"type": "stop"})
            except Exception:
                pass
            
            if self._process and self._process.is_alive():
                self._process.terminate()
                
                # Use executor to join with timeout without blocking the event loop
                loop = asyncio.get_running_loop()
                try:
                    await loop.run_in_executor(None, self._process.join, timeout)
                except Exception:
                    pass
                
                if self._process.is_alive():
                    logger.warning(f"Session {self.id} worker (PID {self._process.pid}) did not exit after {timeout}s; killing.")
                    try:
                        self._process.kill()
                        await loop.run_in_executor(None, self._process.join)
                    except Exception as e:
                        logger.error(f"Failed to kill session {self.id} worker: {e}")
            
            self.status = "stopped"
            if self._listener_task:
                self._listener_task.cancel()
                try:
                    await self._listener_task
                except asyncio.CancelledError:
                    pass
                except Exception:
                    pass

    def get_info(self) -> SessionInfo:
        return SessionInfo(
            id=self.id,
            status=self.status,
            created_at=self.created_at,
            pid=self.pid
        )

# Use class-level list to track all managers for cleanup
_all_managers: List[SessionManager] = []
_atexit_registered = False

def _global_shutdown():
    """Final emergency cleanup for all SessionManagers."""
    for manager in _all_managers:
        manager._shutdown()

class SessionManager:
    def __init__(self):
        self._sessions: Dict[str, StataSession] = {}
        self._default_session_id = "default"
        _all_managers.append(self)
        global _atexit_registered
        if not _atexit_registered:
            atexit.register(_global_shutdown)
            _atexit_registered = True

    def _shutdown(self) -> None:
        """Emergency cleanup for atexit."""
        for session in list(self._sessions.values()):
            try:
                if session._process and session._process.is_alive():
                    # Be very aggressive in atexit, we don't have much time
                    session._process.kill()
                    session._process.join(timeout=0.1)
            except Exception:
                pass
            try:
                session._parent_conn.close()
            except Exception:
                pass
        self._sessions.clear()

    async def start(self):
        # Start default session
        await self.get_or_create_session(self._default_session_id)

    async def get_or_create_session(self, session_id: str) -> StataSession:
        if session_id not in self._sessions:
            logger.info(f"Creating new Stata session: {session_id}")
            session = StataSession(session_id)
            self._sessions[session_id] = session
            # Give it more time to start up on CI (especially Stata's first init)
            # but don't wait if the process dies or status changes.
            timeout = 30.0
            start_time = asyncio.get_running_loop().time()
            while session.status == "starting" and asyncio.get_running_loop().time() - start_time < timeout:
                if not session._process.is_alive():
                    # Process died before reaching ready
                    session.status = "error"
                    break
                await asyncio.sleep(0.1)
                
        return self._sessions[session_id]

    def get_session(self, session_id: str) -> StataSession:
        if session_id not in self._sessions:
            raise ValueError(f"Session {session_id} not found.")
        return self._sessions[session_id]

    def list_sessions(self) -> List[SessionInfo]:
        return [s.get_info() for s in self._sessions.values()]

    async def stop_session(self, session_id: str):
        if session_id in self._sessions:
            await self._sessions[session_id].stop()
            del self._sessions[session_id]

    async def stop_all(self):
        tasks = [s.stop() for s in self._sessions.values()]
        if tasks:
            await asyncio.gather(*tasks)
        self._sessions.clear()
