"""
Server State Manager - Reference counting for server lifecycle.

Manages the CodeContext server lifecycle using reference counting to handle
multiple concurrent Claude Code sessions safely. Each session registers
itself on start and unregisters on end; server stops when count reaches 0.

State file location: ~/.codecontext/server_state.json
"""

import json
import os
import sys
import subprocess
import time
from dataclasses import dataclass, field, asdict
from datetime import datetime, timedelta
from pathlib import Path
from typing import Dict, Optional
from contextlib import contextmanager

# Add codecontext to path for imports (hooks/ is inside codecontext/)
sys.path.insert(0, str(Path(__file__).parent.parent))

from cli_utils import CODECONTEXT_ROOT


# State file location
SERVER_STATE_FILE = CODECONTEXT_ROOT / "server_state.json"
SERVER_STATE_LOCK = CODECONTEXT_ROOT / "server_state.lock"

# Server startup settings
SERVER_STARTUP_TIMEOUT = 60  # seconds to wait for server health
SERVER_HEALTH_CHECK_INTERVAL = 0.5  # seconds between health checks


@contextmanager
def file_lock(lock_path: Path, timeout: float = 10.0):
    """
    Simple cross-platform file lock using atomic file creation.

    Uses lock file creation/removal for mutual exclusion.
    Not bulletproof but sufficient for our use case.
    """
    lock_path = Path(lock_path)
    lock_path.parent.mkdir(parents=True, exist_ok=True)

    start_time = time.time()
    acquired = False

    while time.time() - start_time < timeout:
        try:
            # Try to create lock file exclusively
            fd = os.open(str(lock_path), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
            os.write(fd, str(os.getpid()).encode())
            os.close(fd)
            acquired = True
            break
        except FileExistsError:
            # Lock exists - check if stale (older than 30 seconds)
            try:
                if lock_path.exists():
                    age = time.time() - lock_path.stat().st_mtime
                    if age > 30:
                        # Stale lock - remove and retry
                        lock_path.unlink()
                        continue
            except OSError:
                pass
            time.sleep(0.1)
        except OSError:
            time.sleep(0.1)

    try:
        yield acquired
    finally:
        if acquired:
            try:
                lock_path.unlink()
            except OSError:
                pass


@dataclass
class ServerState:
    """
    Tracks server PID and active session references.

    Attributes:
        server_pid: PID of the running server process, or None
        started_at: ISO timestamp when server was started
        active_sessions: Dict mapping session_id -> registration timestamp
    """
    server_pid: Optional[int] = None
    started_at: Optional[str] = None
    active_sessions: Dict[str, str] = field(default_factory=dict)

    def register_session(self, session_id: str) -> None:
        """Register a session as active."""
        if session_id:
            self.active_sessions[session_id] = datetime.utcnow().isoformat() + "Z"

    def unregister_session(self, session_id: str) -> None:
        """Unregister a session."""
        if session_id and session_id in self.active_sessions:
            del self.active_sessions[session_id]

    def get_active_count(self) -> int:
        """Get count of active sessions."""
        return len(self.active_sessions)

    def cleanup_stale_sessions(self, max_age_hours: int = 24) -> int:
        """
        Remove sessions older than max_age_hours.

        This handles crashed/abandoned sessions that didn't properly unregister.

        Returns:
            Number of stale sessions removed
        """
        cutoff = datetime.utcnow() - timedelta(hours=max_age_hours)
        stale_sessions = []

        for session_id, registered_at in self.active_sessions.items():
            try:
                # Parse ISO timestamp
                reg_time = datetime.fromisoformat(registered_at.replace("Z", "+00:00"))
                # Remove timezone for comparison
                reg_time = reg_time.replace(tzinfo=None)
                if reg_time < cutoff:
                    stale_sessions.append(session_id)
            except (ValueError, TypeError):
                # Invalid timestamp - mark for cleanup
                stale_sessions.append(session_id)

        for session_id in stale_sessions:
            del self.active_sessions[session_id]

        return len(stale_sessions)

    def is_server_running(self) -> bool:
        """
        Check if server process is actually alive.

        Uses os.kill(pid, 0) which doesn't send a signal but checks
        if the process exists and we have permission to signal it.
        """
        if not self.server_pid:
            return False

        try:
            # Signal 0 doesn't kill - just checks if process exists
            os.kill(self.server_pid, 0)
            return True
        except OSError:
            # Process doesn't exist or no permission
            return False
        except Exception:
            return False

    def save(self) -> None:
        """
        Save state to file with file locking.

        Uses file locking to prevent race conditions between
        concurrent sessions trying to update state simultaneously.
        """
        CODECONTEXT_ROOT.mkdir(parents=True, exist_ok=True)

        with file_lock(SERVER_STATE_LOCK, timeout=10) as acquired:
            # Save regardless of lock acquisition (better than losing state)
            SERVER_STATE_FILE.write_text(json.dumps(asdict(self), indent=2))

    @classmethod
    def load(cls) -> "ServerState":
        """
        Load state from file with file locking.

        Returns fresh ServerState if file doesn't exist or is corrupted.
        """
        with file_lock(SERVER_STATE_LOCK, timeout=10):
            try:
                if SERVER_STATE_FILE.exists():
                    data = json.loads(SERVER_STATE_FILE.read_text())
                    return cls(
                        server_pid=data.get("server_pid"),
                        started_at=data.get("started_at"),
                        active_sessions=data.get("active_sessions", {})
                    )
            except (json.JSONDecodeError, IOError):
                # Corrupted or inaccessible - return fresh state
                pass

        return cls()


def cleanup_zombie_servers(port_range: range = range(8765, 8775)) -> int:
    """
    Kill any processes listening on the CodeContext port range.

    This prevents the "port busy" issue where zombie servers cause
    new servers to start on different ports.

    Args:
        port_range: Range of ports to check (default: 8765-8774)

    Returns:
        Number of processes killed
    """
    killed = 0

    if sys.platform == "win32":
        try:
            # Get netstat output once
            result = subprocess.run(
                ["netstat", "-ano"],
                capture_output=True,
                text=True,
                timeout=5
            )

            # Find PIDs listening on our ports
            pids_to_kill = set()
            for line in result.stdout.splitlines():
                for port in port_range:
                    if f":{port}" in line and "LISTENING" in line:
                        parts = line.split()
                        if parts:
                            pids_to_kill.add(parts[-1])

            # Kill each PID
            for pid in pids_to_kill:
                try:
                    subprocess.run(
                        ["taskkill", "/PID", pid, "/F"],
                        capture_output=True,
                        timeout=5
                    )
                    killed += 1
                except Exception:
                    pass

        except Exception:
            pass
    else:
        # Unix: use lsof to find processes
        for port in port_range:
            try:
                result = subprocess.run(
                    ["lsof", "-ti", f":{port}"],
                    capture_output=True,
                    text=True,
                    timeout=5
                )
                for pid in result.stdout.strip().split():
                    try:
                        os.kill(int(pid), 9)  # SIGKILL
                        killed += 1
                    except Exception:
                        pass
            except Exception:
                pass

    return killed


def start_server_background() -> int:
    """
    Start the CodeContext server as a detached background process.

    Cleans up any zombie servers on the port range first.

    Returns:
        PID of the started server process

    Raises:
        RuntimeError: If server fails to start
    """
    # Clean up any zombie servers first
    killed = cleanup_zombie_servers()
    if killed > 0:
        # Give OS time to release the ports
        time.sleep(1)

    # Find server.py path
    server_path = Path(__file__).parent.parent / "server.py"

    if not server_path.exists():
        raise RuntimeError(f"Server script not found: {server_path}")

    # Get Python executable
    python_exe = sys.executable

    # Windows-specific: use CREATE_NO_WINDOW to run detached
    if sys.platform == "win32":
        CREATE_NO_WINDOW = 0x08000000
        CREATE_NEW_PROCESS_GROUP = 0x00000200
        DETACHED_PROCESS = 0x00000008

        process = subprocess.Popen(
            [python_exe, str(server_path)],
            cwd=str(server_path.parent),
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            stdin=subprocess.DEVNULL,
            creationflags=CREATE_NO_WINDOW | CREATE_NEW_PROCESS_GROUP | DETACHED_PROCESS
        )
    else:
        # Unix: use start_new_session to detach
        process = subprocess.Popen(
            [python_exe, str(server_path)],
            cwd=str(server_path.parent),
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            stdin=subprocess.DEVNULL,
            start_new_session=True
        )

    return process.pid


def shutdown_server(pid: int) -> bool:
    """
    Gracefully shutdown server by PID.

    On Windows, uses taskkill. On Unix, sends SIGTERM.

    Args:
        pid: Process ID of server to shutdown

    Returns:
        True if shutdown was successful (or process already dead),
        False if shutdown failed
    """
    try:
        if sys.platform == "win32":
            # Windows: use taskkill for graceful shutdown
            subprocess.run(
                ["taskkill", "/PID", str(pid), "/F"],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                timeout=10
            )
        else:
            # Unix: send SIGTERM for graceful shutdown
            import signal
            os.kill(pid, signal.SIGTERM)

            # Give process time to shutdown
            for _ in range(10):
                time.sleep(0.5)
                try:
                    os.kill(pid, 0)
                except OSError:
                    break  # Process is dead
            else:
                # Process still alive - force kill
                os.kill(pid, signal.SIGKILL)

        return True

    except subprocess.TimeoutExpired:
        # Taskkill timed out
        return False
    except ProcessLookupError:
        # Already dead - success
        return True
    except Exception:
        return False


def wait_for_server_health(base_url: str = "http://localhost:8765", timeout: int = None) -> bool:
    """
    Wait for server to become healthy.

    Args:
        base_url: Server base URL
        timeout: Maximum seconds to wait (default: SERVER_STARTUP_TIMEOUT)

    Returns:
        True if server became healthy, False if timeout
    """
    import urllib.request
    import urllib.error

    if timeout is None:
        timeout = SERVER_STARTUP_TIMEOUT

    health_url = f"{base_url}/health"
    deadline = time.time() + timeout

    while time.time() < deadline:
        try:
            req = urllib.request.Request(health_url)
            with urllib.request.urlopen(req, timeout=2) as response:
                if response.status == 200:
                    return True
        except (urllib.error.URLError, urllib.error.HTTPError, OSError):
            pass

        time.sleep(SERVER_HEALTH_CHECK_INTERVAL)

    return False


def ensure_server_running(base_url: str = "http://localhost:8765") -> bool:
    """
    Ensure server is running, starting it if necessary.

    This is the main entry point for CLI commands that need the server.

    Args:
        base_url: Server base URL for health checks

    Returns:
        True if server is running (or was started successfully),
        False if server could not be started
    """
    import urllib.request
    import urllib.error

    # Quick health check first
    try:
        req = urllib.request.Request(f"{base_url}/health")
        with urllib.request.urlopen(req, timeout=2) as response:
            if response.status == 200:
                return True
    except (urllib.error.URLError, urllib.error.HTTPError, OSError):
        pass

    # Server not responding - check state and start if needed
    state = ServerState.load()

    # Clean up if PID is stale
    if state.server_pid and not state.is_server_running():
        state.server_pid = None
        state.started_at = None
        state.save()

    # Start server if not running
    if not state.server_pid:
        try:
            pid = start_server_background()
            state.server_pid = pid
            state.started_at = datetime.utcnow().isoformat() + "Z"
            state.save()
        except Exception as e:
            print(f"[CodeContext] Failed to start server: {e}", file=sys.stderr)
            return False

    # Wait for server to become healthy
    return wait_for_server_health(base_url)
