"""
server/lsp_server.py — Lightweight LSP implementation for RPy.
"""

from __future__ import annotations
import json
import socketserver
import sys
from typing import Any, Dict, List, Optional
from transpiler.middleend.query_service import query_service
import os

class ProfileLoader:
    _cache: Dict[str, Dict[str, Any]] = {}

    @classmethod
    def get_profile(cls, root: str) -> Dict[str, Any]:
        profile_path = os.path.join(root, ".rpy_profile")
        if not os.path.exists(profile_path):
            return {}
        
        # Simple polling/caching
        try:
            mtime = os.path.getmtime(profile_path)
            cached = cls._cache.get(profile_path)
            if cached and cached.get("mtime") == mtime:
                return cached["data"]
            
            with open(profile_path, "r") as f:
                data = json.load(f)
            cls._cache[profile_path] = {"mtime": mtime, "data": data}
            return data
        except Exception:
            return {}

class StdioLSPHandler:
    document_cache: Dict[str, str] = {}
    running: bool = True

    def handle(self):
        stdin = sys.stdin.buffer
        while self.running:
            try:
                line = stdin.readline()
                if not line:
                    break
                line_str = line.decode("utf-8")
                if not line_str.startswith("Content-Length:"):
                    continue
                
                length = int(line_str.split(":")[1].strip())
                # Read the empty line after headers
                stdin.readline()
                
                content = stdin.read(length).decode("utf-8")
                if not content:
                    break
                
                request = json.loads(content)
                self.process_request(request)
            except EOFError:
                break
            except Exception as e:
                import traceback
                with open("lsp_crash.log", "a") as f:
                    f.write(f"Error in handle loop: {e}\n{traceback.format_exc()}\n")
                break

    def process_request(self, request: Dict[str, Any]):
        method = request.get("method")
        params = request.get("params", {})
        request_id = request.get("id")

        if method == "initialize":
            self.handle_initialize(request_id, params)
        elif method == "initialized":
            pass # No-op for now
        elif method == "textDocument/didOpen":
            self.handle_sync(params)
        elif method == "textDocument/didChange":
            self.handle_sync(params)
        elif method == "textDocument/hover":
            self.handle_hover(request_id, params)
        elif method == "textDocument/completion":
            self.handle_completion(request_id, params)
        elif method == "textDocument/definition":
            self.handle_definition(request_id, params)
        elif method == "textDocument/references":
            self.handle_references(request_id, params)
        elif method == "textDocument/codeAction":
            self.handle_code_action(request_id, params)
        elif method == "shutdown":
            self.handle_shutdown(request_id)
        elif method == "exit":
            self.handle_exit()

    def handle_initialize(self, request_id: Any, params: Dict[str, Any]):
        root_uri = params.get("rootUri")
        if root_uri and root_uri.startswith("file://"):
            root_path = root_uri.replace("file://", "")
            if os.name == "nt" and root_path.startswith("/"):
                root_path = root_path[1:]
            query_service.set_workspace_root(root_path)

        self.send_response(request_id, {
            "capabilities": {
                "textDocumentSync": 1, 
                "hoverProvider": True,
                "definitionProvider": True,
                "completionProvider": {
                    "resolveProvider": True,
                    "triggerCharacters": [".", ":"]
                },
                "codeActionProvider": True
            },
            "serverInfo": {
                "name": "RPy Language Server",
                "version": "1.0.0"
            }
        })

    def handle_shutdown(self, request_id: Any):
        self.send_response(request_id, None)

    def handle_exit(self):
        self.running = False
        sys.exit(0)

    def handle_code_action(self, request_id: Any, params: Dict[str, Any]):
        uri = params.get("textDocument", {}).get("uri", "")
        diagnostics = params.get("context", {}).get("diagnostics", [])
        actions = []
        
        for diag in diagnostics:
            msg = diag.get("message", "")
            if "Try" in msg or "try" in msg.lower():
                actions.append({
                    "title": "💡 RPy Quick Fix: Rewrite using safe pcall() pattern",
                    "kind": "quickfix",
                    "diagnostics": [diag],
                    "isPreferred": True,
                    "edit": {
                        "changes": {
                            uri: [{
                                "range": diag["range"],
                                "newText": "# RPy doesn't support try/except natively.\n# Wrap your unsafe execution in a target function:\ndef execute_safe():\n    pass # Replace with execution\n\nsuccess, result = pcall(execute_safe)\nif not success:\n    print('Handled error:', result)"
                            }]
                        }
                    }
                })
            elif "yield" in msg.lower():
                actions.append({
                    "title": "💡 RPy Quick Fix: Use task.wait() or Roblox yielding functions",
                    "kind": "quickfix",
                    "diagnostics": [diag],
                    "edit": {
                        "changes": {
                            uri: [{
                                "range": diag["range"],
                                "newText": "import task\ntask.wait(1.0)"
                            }]
                        }
                    }
                })
                
        self.send_response(request_id, actions)

    def handle_sync(self, params: Dict[str, Any]):
        doc = params.get("textDocument", {})
        uri = doc.get("uri", "")
        
        changes = params.get("contentChanges")
        text = changes[0].get("text") if changes else doc.get("text")
            
        if uri:
            affected_uris = query_service.reload_file(uri, text)
            for affected_uri in affected_uris:
                diagnostics = query_service.get_diagnostics(affected_uri)
                diag_list = [
                    {
                        "range": {
                            "start": {"line": d["line"] - 1, "character": d["col"]},
                            "end": {"line": d["line"] - 1, "character": d["col"] + 1}
                        },
                        "severity": {
                            "ERROR": 1,
                            "WARNING": 2,
                            "HINT": 4
                        }.get(d["severity"], 1),
                        "message": d["message"],
                        "source": "RPy",
                        "code": d.get("code", "RPY000")
                    }
                    for d in diagnostics if d.get("line") is not None
                ]
                self.send_notification("textDocument/publishDiagnostics", {
                    "uri": affected_uri,
                    "diagnostics": diag_list
                })

    def handle_hover(self, request_id: Any, params: Dict[str, Any]):
        pos = params.get("position", {})
        uri = params.get("textDocument", {}).get("uri", "")
        line = pos.get("line", 0) + 1
        col = pos.get("character", 0)
        
        hover_text = query_service.get_hover(uri, line, col)
        if hover_text:
            self.send_response(request_id, {
                "contents": {
                    "kind": "markdown",
                    "value": hover_text
                }
            })
            return
        
        self.send_response(request_id, None)

    def handle_completion(self, request_id: Any, params: Dict[str, Any]):
        pos = params.get("position", {})
        uri = params.get("textDocument", {}).get("uri", "")
        line = pos.get("line", 0) + 1
        col = pos.get("character", 0)
        
        items = query_service.get_completions(uri, line, col)
        # Kind mapping done in query_service now, or we can ensure it here
        self.send_response(request_id, {"isIncomplete": False, "items": items})

    def handle_definition(self, request_id: Any, params: Dict[str, Any]):
        pos = params.get("position", {})
        uri = params.get("textDocument", {}).get("uri", "")
        line = pos.get("line", 0) + 1
        col = pos.get("character", 0)
        
        location = query_service.get_definition(uri, line, col)
        if location:
            self.send_response(request_id, location)
            return
        self.send_response(request_id, None)

    def handle_references(self, request_id: Any, params: Dict[str, Any]):
        pos = params.get("position", {})
        uri = params.get("textDocument", {}).get("uri", "")
        line = pos.get("line", 0) + 1
        col = pos.get("character", 0)
        
        refs = query_service.find_references(uri, line, col)
        self.send_response(request_id, refs)

    def send_response(self, request_id: Any, result: Any):
        response = {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": result
        }
        self._send(response)

    def send_notification(self, method: str, params: Any):
        notification = {
            "jsonrpc": "2.0",
            "method": method,
            "params": params
        }
        self._send(notification)

    def _send(self, data: Dict[str, Any]):
        body = json.dumps(data, separators=(',', ':'))
        message = f"Content-Length: {len(body)}\r\n\r\n{body}"
        sys.stdout.buffer.write(message.encode("utf-8"))
        sys.stdout.buffer.flush()

def start_lsp_server(host: Optional[str] = None, port: Optional[int] = None):
    """Entry point for starting the LSP server. Currently supports stdio."""
    handler = StdioLSPHandler()
    handler.handle()

def main() -> int:
    start_lsp_server()
    return 0

if __name__ == "__main__":
    main()
