"""
LSP Client - JSON-RPC 2.0 client for Language Server Protocol

Provides a client for communicating with LSP servers over stdio.

Examples:
    client = LSPClient("pyright-langserver", ["--stdio"])
    client.start()
    client.initialize(root_uri="file:///project")

    # Get hover info
    result = client.hover("file:///app.py", line=10, character=5)

    client.shutdown()
"""

import json
import logging
import subprocess
import threading
from dataclasses import dataclass, field
from pathlib import Path
from queue import Empty, Queue
from typing import Any, Optional

logger = logging.getLogger(__name__)


class LSPClientError(Exception):
    """Base exception for LSP client errors"""

    pass


class LSPTimeoutError(LSPClientError):
    """Raised when a request times out"""

    pass


class LSPRequestError(LSPClientError):
    """Raised when the server returns an error response"""

    def __init__(self, code: int, message: str, data: Any = None):
        self.code = code
        self.message = message
        self.data = data
        super().__init__(f"LSP Error {code}: {message}")


@dataclass
class LSPClientConfig:
    """Configuration for LSP client"""

    command: str
    args: list[str] = field(default_factory=list)
    cwd: Optional[Path] = None
    env: Optional[dict[str, str]] = None
    timeout: float = 30.0
    init_timeout: float = 60.0


class LSPClient:
    """
    JSON-RPC 2.0 client for Language Server Protocol

    Handles communication with LSP servers over stdio using
    the JSON-RPC 2.0 protocol with Content-Length headers.

    Example:
        client = LSPClient("pyright-langserver", ["--stdio"])
        client.start()
        client.initialize(root_uri="file:///project")

        # Send requests
        result = client.send_request("textDocument/hover", {
            "textDocument": {"uri": "file:///app.py"},
            "position": {"line": 10, "character": 5}
        })

        client.shutdown()
    """

    def __init__(
        self,
        command: str,
        args: Optional[list[str]] = None,
        cwd: Optional[Path] = None,
        env: Optional[dict[str, str]] = None,
        timeout: float = 30.0,
        init_timeout: float = 60.0,
    ):
        """
        Initialize LSP client

        Args:
            command: Server command/executable
            args: Command line arguments
            cwd: Working directory for server
            env: Environment variables
            timeout: Default request timeout in seconds
            init_timeout: Timeout for initialize request
        """
        self.config = LSPClientConfig(
            command=command,
            args=args or [],
            cwd=cwd,
            env=env,
            timeout=timeout,
            init_timeout=init_timeout,
        )

        self._process: Optional[subprocess.Popen] = None
        self._request_id = 0
        self._pending_requests: dict[int, Queue] = {}
        self._reader_thread: Optional[threading.Thread] = None
        self._running = False
        self._lock = threading.Lock()
        self._initialized = False
        self._server_capabilities: dict[str, Any] = {}

    @property
    def is_running(self) -> bool:
        """Check if client is running"""
        return self._running and self._process is not None

    @property
    def is_initialized(self) -> bool:
        """Check if server is initialized"""
        return self._initialized

    @property
    def server_capabilities(self) -> dict[str, Any]:
        """Get server capabilities from initialize response"""
        return self._server_capabilities.copy()

    def start(self) -> None:
        """
        Start the LSP server process

        Raises:
            LSPClientError: If server fails to start
        """
        if self._running:
            return

        cmd = [self.config.command] + self.config.args

        try:
            self._process = subprocess.Popen(
                cmd,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                cwd=self.config.cwd,
                env=self.config.env,
            )
        except FileNotFoundError:
            raise LSPClientError(f"Server command not found: {self.config.command}")
        except Exception as e:
            raise LSPClientError(f"Failed to start server: {e}")

        self._running = True

        # Start reader thread
        self._reader_thread = threading.Thread(target=self._read_responses, daemon=True)
        self._reader_thread.start()

        logger.info(f"Started LSP server: {' '.join(cmd)}")

    def stop(self) -> None:
        """
        Stop the LSP server process

        Tries graceful shutdown first, then force kills if needed.
        """
        if not self._running:
            return

        self._running = False

        if self._process:
            try:
                # Try graceful termination
                self._process.terminate()
                self._process.wait(timeout=5.0)
            except subprocess.TimeoutExpired:
                # Force kill
                self._process.kill()
                self._process.wait()
            finally:
                self._process = None

        self._initialized = False
        self._server_capabilities = {}

        logger.info("Stopped LSP server")

    def restart(self) -> None:
        """
        Restart the LSP server

        Stops the current server and starts a new one.
        Does NOT re-initialize - call initialize() after restart.
        """
        self.stop()
        self.start()

    def initialize(
        self,
        root_uri: str,
        capabilities: Optional[dict[str, Any]] = None,
        workspace_folders: Optional[list[dict[str, str]]] = None,
        timeout: Optional[float] = None,
    ) -> dict[str, Any]:
        """
        Initialize the LSP server

        Args:
            root_uri: Root URI of the workspace (file:///path/to/project)
            capabilities: Client capabilities
            workspace_folders: List of workspace folders
            timeout: Timeout for initialize request (uses init_timeout if None)

        Returns:
            Initialize result with server capabilities

        Raises:
            LSPClientError: If initialization fails
        """
        if not self.is_running:
            raise LSPClientError("Server not running")

        if self._initialized:
            return {"capabilities": self._server_capabilities}

        # Default client capabilities
        default_capabilities = {
            "textDocument": {
                "hover": {"contentFormat": ["markdown", "plaintext"]},
                "definition": {"linkSupport": True},
                "references": {},
                "documentSymbol": {"hierarchicalDocumentSymbolSupport": True},
            },
            "workspace": {
                "workspaceFolders": True,
            },
        }

        params = {
            "processId": None,
            "rootUri": root_uri,
            "capabilities": capabilities or default_capabilities,
            "initializationOptions": {},
        }

        if workspace_folders:
            params["workspaceFolders"] = workspace_folders

        try:
            result = self.send_request(
                "initialize",
                params,
                timeout=timeout or self.config.init_timeout,
            )

            self._server_capabilities = result.get("capabilities", {})
            self._initialized = True

            # Send initialized notification
            self.send_notification("initialized", {})

            logger.info("LSP server initialized")
            return result

        except Exception as e:
            raise LSPClientError(f"Initialize failed: {e}")

    def shutdown(self) -> None:
        """
        Shutdown the LSP server gracefully

        Sends shutdown request followed by exit notification.
        """
        if not self.is_running:
            return

        try:
            if self._initialized:
                self.send_request("shutdown", None, timeout=10.0)
                self.send_notification("exit", None)
        except Exception as e:
            logger.warning(f"Error during shutdown: {e}")
        finally:
            self.stop()

    def send_request(
        self,
        method: str,
        params: Any,
        timeout: Optional[float] = None,
    ) -> Any:
        """
        Send a request and wait for response

        Args:
            method: LSP method name
            params: Request parameters
            timeout: Request timeout (uses default if None)

        Returns:
            Response result

        Raises:
            LSPTimeoutError: If request times out
            LSPRequestError: If server returns error
            LSPClientError: For other errors
        """
        if not self.is_running:
            raise LSPClientError("Server not running")

        timeout = timeout or self.config.timeout

        with self._lock:
            self._request_id += 1
            request_id = self._request_id
            response_queue: Queue = Queue()
            self._pending_requests[request_id] = response_queue

        message = {
            "jsonrpc": "2.0",
            "id": request_id,
            "method": method,
            "params": params,
        }

        try:
            self._send_message(message)

            # Wait for response
            try:
                response = response_queue.get(timeout=timeout)
            except Empty:
                raise LSPTimeoutError(f"Request timed out: {method}")

            # Check for error response
            if "error" in response:
                error = response["error"]
                raise LSPRequestError(
                    code=error.get("code", -1),
                    message=error.get("message", "Unknown error"),
                    data=error.get("data"),
                )

            return response.get("result")

        finally:
            with self._lock:
                self._pending_requests.pop(request_id, None)

    def send_notification(self, method: str, params: Any) -> None:
        """
        Send a notification (no response expected)

        Args:
            method: LSP method name
            params: Notification parameters
        """
        if not self.is_running:
            raise LSPClientError("Server not running")

        message = {
            "jsonrpc": "2.0",
            "method": method,
            "params": params,
        }

        self._send_message(message)

    def _send_message(self, message: dict[str, Any]) -> None:
        """Send a JSON-RPC message to the server"""
        if not self._process or not self._process.stdin:
            raise LSPClientError("Server stdin not available")

        content = json.dumps(message)
        content_bytes = content.encode("utf-8")

        header = f"Content-Length: {len(content_bytes)}\r\n\r\n"
        full_message = header.encode("utf-8") + content_bytes

        try:
            self._process.stdin.write(full_message)
            self._process.stdin.flush()
            logger.debug(f"Sent: {message.get('method', 'response')}")
        except Exception as e:
            raise LSPClientError(f"Failed to send message: {e}")

    def _read_responses(self) -> None:
        """Background thread that reads responses from server"""
        while self._running and self._process and self._process.stdout:
            try:
                message = self._read_message()
                if message is None:
                    break

                self._handle_message(message)

            except Exception as e:
                if self._running:
                    logger.error(f"Error reading response: {e}")
                break

    def _read_message(self) -> Optional[dict[str, Any]]:
        """Read a single JSON-RPC message from server"""
        if not self._process or not self._process.stdout:
            return None

        # Read headers
        headers: dict[str, str] = {}
        while True:
            line = self._process.stdout.readline()
            if not line:
                return None

            line_str = line.decode("utf-8").strip()
            if not line_str:
                break

            if ":" in line_str:
                key, value = line_str.split(":", 1)
                headers[key.strip()] = value.strip()

        # Get content length
        content_length = int(headers.get("Content-Length", 0))
        if content_length == 0:
            return None

        # Read content
        content = self._process.stdout.read(content_length)
        if not content:
            return None

        try:
            return json.loads(content.decode("utf-8"))
        except json.JSONDecodeError as e:
            logger.error(f"Invalid JSON response: {e}")
            return None

    def _handle_message(self, message: dict[str, Any]) -> None:
        """Handle an incoming message"""
        if "id" in message:
            # Response to a request
            request_id = message["id"]
            with self._lock:
                if request_id in self._pending_requests:
                    self._pending_requests[request_id].put(message)
                else:
                    logger.warning(f"Unknown request id: {request_id}")
        else:
            # Notification from server
            method = message.get("method", "")
            logger.debug(f"Received notification: {method}")

    # Convenience methods for common LSP operations

    def hover(self, uri: str, line: int, character: int) -> Optional[dict[str, Any]]:
        """
        Get hover information at position

        Args:
            uri: Document URI (file:///path)
            line: Line number (0-based)
            character: Character offset (0-based)

        Returns:
            Hover result or None
        """
        return self.send_request(
            "textDocument/hover",
            {
                "textDocument": {"uri": uri},
                "position": {"line": line, "character": character},
            },
        )

    def definition(self, uri: str, line: int, character: int) -> Optional[list[dict[str, Any]]]:
        """
        Go to definition

        Args:
            uri: Document URI
            line: Line number (0-based)
            character: Character offset (0-based)

        Returns:
            List of locations or None
        """
        result = self.send_request(
            "textDocument/definition",
            {
                "textDocument": {"uri": uri},
                "position": {"line": line, "character": character},
            },
        )

        # Normalize to list
        if result is None:
            return None
        if isinstance(result, dict):
            return [result]
        return result

    def references(
        self,
        uri: str,
        line: int,
        character: int,
        include_declaration: bool = True,
    ) -> Optional[list[dict[str, Any]]]:
        """
        Find references

        Args:
            uri: Document URI
            line: Line number (0-based)
            character: Character offset (0-based)
            include_declaration: Include declaration in results

        Returns:
            List of locations or None
        """
        return self.send_request(
            "textDocument/references",
            {
                "textDocument": {"uri": uri},
                "position": {"line": line, "character": character},
                "context": {"includeDeclaration": include_declaration},
            },
        )

    def document_symbols(self, uri: str) -> Optional[list[dict[str, Any]]]:
        """
        Get document symbols

        Args:
            uri: Document URI

        Returns:
            List of symbols or None
        """
        return self.send_request(
            "textDocument/documentSymbol",
            {"textDocument": {"uri": uri}},
        )

    def did_open(self, uri: str, language_id: str, text: str, version: int = 1) -> None:
        """
        Notify server that a document was opened

        Args:
            uri: Document URI
            language_id: Language identifier (python, javascript, etc.)
            text: Document content
            version: Document version
        """
        self.send_notification(
            "textDocument/didOpen",
            {
                "textDocument": {
                    "uri": uri,
                    "languageId": language_id,
                    "version": version,
                    "text": text,
                }
            },
        )

    def did_close(self, uri: str) -> None:
        """
        Notify server that a document was closed

        Args:
            uri: Document URI
        """
        self.send_notification(
            "textDocument/didClose",
            {"textDocument": {"uri": uri}},
        )

    def did_change(
        self,
        uri: str,
        version: int,
        changes: list[dict[str, Any]],
    ) -> None:
        """
        Notify server of document changes

        Args:
            uri: Document URI
            version: New document version
            changes: List of content changes
        """
        self.send_notification(
            "textDocument/didChange",
            {
                "textDocument": {"uri": uri, "version": version},
                "contentChanges": changes,
            },
        )

    def __enter__(self) -> "LSPClient":
        """Context manager entry"""
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        """Context manager exit"""
        self.shutdown()
