"""
transpiler/query_service.py — High-level semantic queries for LSP and IDEs.
"""

from __future__ import annotations
import ast
from typing import List, Optional, Any, Dict
from pathlib import Path

from transpiler.infrastructure.dependency_graph import DependencyGraph
from transpiler.infrastructure.project_context import ProjectContext
from transpiler.infrastructure.visitor import IRVisitor
from transpiler.middleend import ir
from transpiler.middleend.analyzer import SemanticAnalyzer
from transpiler.optimization.ir_transformer import PythonToIRTransformer
from transpiler.infrastructure.diagnostics import DiagnosticManager
from transpiler.infrastructure.linter import lint_ir

def _query_transform(tree: ast.Module, filename: str = "<string>", project: Optional[ProjectContext] = None, module_name: str = "main", uri: Optional[str] = None, diagnostics: Optional[DiagnosticManager] = None):
    """Lightweight transform for IDE queries: parse -> IR only, no optimizer.
    This preserves original lineno/col_offset on every IRNode so NodeFinder can
    locate symbols at precise cursor positions.
    """
    transformer = PythonToIRTransformer()
    ir_module = transformer.transform(tree)
    # Semantic analysis populates inferred_type and definition_location on all nodes
    analyzer = SemanticAnalyzer(diagnostics=diagnostics, project=project, module_name=module_name, uri=uri)
    analyzer.analyze(ir_module)
    return ir_module

class NodeFinder(IRVisitor):
    """Finds the most specific IR node at a given line and column."""
    def __init__(self, line: int, col: int):
        self.line = line
        self.col = col
        self.found_node: Optional[ir.IRNode] = None

    def _get_node_end(self, node: ir.IRNode) -> int:
        if isinstance(node, ir.IRVariable):
            return node.col_offset + len(node.name)
        if isinstance(node, (ir.IRFunctionDef, ir.IRClassDef)):
            # Heuristic: 'def name(' or 'class name('
            return node.col_offset + len(node.name) + 6
        if isinstance(node, ir.IRPropertyAccess):
            return self._get_node_end(node.receiver) + 1 + len(node.property)
        if isinstance(node, ir.IRLiteral):
            return node.col_offset + len(str(node.value))
        if isinstance(node, (ir.IRImport, ir.IRImportFrom)):
            # Heuristic: the whole line for imports
            return node.col_offset + 100
        return node.col_offset + 1 # Default

    def visit(self, node: Any) -> Any:
        # Narrow the type of node to IRNode for col_offset/lineno access
        if isinstance(node, ir.IRNode):
            if getattr(node, "lineno", 0) == self.line:
                start = node.col_offset
                end = self._get_node_end(node)
                
                # Check if cursor is within bounds
                if start <= self.col < end:
                    # We prefer the smallest (narrowest) node that contains the cursor
                    current = self.found_node
                    if current is None:
                        self.found_node = node
                    elif isinstance(current, ir.IRNode):
                        old_width = getattr(current, "end_col", self._get_node_end(current)) - current.col_offset
                        new_width = end - start
                        
                        if new_width <= old_width:
                            self.found_node = node
                            
        # Always traverse children
        self.generic_visit(node)
        return super().visit(node)

class QueryService:
    def __init__(self):
        self.project_context = ProjectContext()
        self.dep_graph: Optional[DependencyGraph] = None
        self.workspace_root: Optional[Path] = None
        
        # Cache per file: {uri: {"module": IRModule, "diagnostics": [diag], "text": str}}
        self._cache: Dict[str, Dict[str, Any]] = {}
        self.diagnostics_manager = DiagnosticManager(silent=True)

    def _normalize_uri(self, uri: str) -> str:
        # Normalize Windows drive letters and separators
        if not uri: return uri
        u = uri.replace("\\", "/")
        if "file://" in u:
            # RPy LSP URI format: file:///C:/... or file://C:/...
            path_part = u.split("file://", 1)[1]
            if path_part.startswith("///"): path_part = path_part[3:]
            elif path_part.startswith("/"): path_part = path_part[1:]
            
            # Re-normalize drive letter to uppercase
            if len(path_part) > 1 and path_part[1] == ":" and path_part[0].islower():
                path_part = path_part[0].upper() + path_part[1:]
            
            return f"file:///{path_part}"
        return u

    def set_workspace_root(self, root_path: str):
        # Normalize path casing for Windows (Drive letter uppercase)
        abs_root = Path(root_path).resolve()
        path_str = abs_root.as_posix()
        if len(path_str) > 1 and path_str[1] == ":" and path_str[0].islower():
            path_str = path_str[0].upper() + path_str[1:]
        self.workspace_root = Path(path_str)
        if self.project_context:
            self.project_context.workspace_root = str(self.workspace_root)
        
        self.dep_graph = DependencyGraph(str(self.workspace_root))
        self.dep_graph.scan_project(ignore_roots=True) # LSP handles many files
        
        # Initial analysis of all files in the graph in topological order
        try:
            build_order = self.dep_graph.get_build_order()
            for file_path in build_order:
                uri = f"file://{file_path.as_posix()}"
                self.reload_file(uri)
        except Exception:
            # Fallback to unordered if cycle exists
            for file_path in self.dep_graph.graph:
                uri = f"file://{file_path.as_posix()}"
                self.reload_file(uri)

    def _get_module_name(self, path: Path) -> str:
        if not self.workspace_root:
            return "main"
            
        def _normalize_p(p: Path) -> str:
            s = p.resolve().as_posix()
            if len(s) > 1 and s[1] == ":" and s[0].islower():
                s = s[0].upper() + s[1:]
            return s

        try:
            norm_path = _normalize_p(path)
            norm_root = _normalize_p(self.workspace_root)
            
            # Ensure path starts with root
            if not norm_path.startswith(norm_root):
                return "main"
                
            rel_str = norm_path[len(norm_root):].lstrip("/")
            rel_path = Path(rel_str)
            parts = list(rel_path.with_suffix("").parts)
            if parts and parts[-1] == "__init__":
                parts.pop()
            return ".".join(parts) if parts else "main"
        except Exception:
            return "main"

    def reload_file(self, uri: str, text: Optional[str] = None) -> List[str]:
        """
        Reloads a file, updates the graph, and re-analyzes affected modules.
        Returns a list of URIs that need diagnostic updates.
        """
        uri = self._normalize_uri(uri)
        path = Path(uri.replace("file:///", "").replace("file://", ""))
        if not text:
            try:
                text = path.read_text(encoding="utf-8")
            except Exception:
                return []

        module_name = self._get_module_name(path)
        
        if self.dep_graph:
            self.dep_graph._process_file(path)

        # Clear diagnostics for this file
        self.diagnostics_manager.clear()
        
        try:
            tree = ast.parse(text)
            ir_module = _query_transform(tree, filename=str(path), project=self.project_context, module_name=module_name, uri=uri, diagnostics=self.diagnostics_manager)
            
            # Run linting
            lint_ir(ir_module, diagnostics=self.diagnostics_manager)
            
            diags = [
                {
                    "line": d.line,
                    "col": d.col,
                    "severity": d.severity.name,
                    "message": d.message,
                    "code": d.code if hasattr(d, "code") else "RPY000"
                }
                for d in self.diagnostics_manager.diagnostics
            ]
            
            self._cache[uri] = {
                "module": ir_module,
                "diagnostics": diags,
                "text": text
            }
        except Exception as e:
            self._cache[uri] = {
                "module": None,
                "diagnostics": [{"line": 1, "col": 0, "severity": "ERROR", "message": f"Parse Error: {e}"}],
                "text": text
            }

        return [uri]

    def get_diagnostics(self, uri: str) -> List[Dict[str, Any]]:
        uri = self._normalize_uri(uri)
        cached = self._cache.get(uri)
        if cached:
            return cached["diagnostics"]
        return []

    def get_hover(self, uri: str, line: int, col: int) -> Optional[str]:
        uri = self._normalize_uri(uri)
        cached = self._cache.get(uri)
        if not cached or not cached["module"]:
            return None
        
        ir_module = cached["module"]
        finder = NodeFinder(line, col)
        finder.visit(ir_module)
        
        node = finder.found_node
        if node:
            if isinstance(node, ir.IRPropertyAccess):
                return f"```python\n(property) {node.property}: {node.inferred_type or 'any'}\n```"
            if isinstance(node, ir.IRMethodCall):
                return f"```python\n(method) {node.method}: {node.inferred_type or 'any'}\n```"
            if isinstance(node, ir.IRVariable):
                return f"```python\n({ 'variable' if not node.inferred_type or not node.inferred_type.startswith('Class:') else 'class' }) {node.name}: {node.inferred_type or 'any'}\n```"
            if isinstance(node, ir.IRExpression) and node.inferred_type:
                return f"```python\n(type) {node.inferred_type}\n```"
            return f"```python\n({type(node).__name__})\n```"
        return None

    def get_completions(self, uri: str, line: int, col: int) -> List[Dict[str, Any]]:
        uri = self._normalize_uri(uri)
        cached = self._cache.get(uri)
        if not cached: return []
        
        source = cached["text"]
        source_lines = source.splitlines()
        if line - 1 >= len(source_lines): return []
        
        curr_line = source_lines[line-1]
        prefix = curr_line[:col]
        
        if prefix.endswith("."):
            receiver_expr = prefix[:-1]
            try:
                import re
                match = re.search(r"([a-zA-Z_0-9\.]+)$", receiver_expr)
                if not match: return []
                
                expr = match.group(0)
                # Lightweight fallback for completions
                context_source = "\n".join(source_lines[:line-1]) + f"\n_tmp = {expr}"
                tree = ast.parse(context_source)
                analyzer = SemanticAnalyzer(diagnostics=self.diagnostics_manager, project=self.project_context, uri=uri)
                ir_module = PythonToIRTransformer().transform(tree)
                analyzer.analyze(ir_module)
                
                last_stmt = ir_module.body[-1]
                if isinstance(last_stmt, ir.IRAssignment):
                    rtype = analyzer._infer_type(last_stmt.value)
                    if rtype:
                        # 1. Check local classes
                        info = analyzer._get_class_info(rtype)
                        # 2. Check SDK types
                        from transpiler.infrastructure.sdk_metadata import get_members
                        sdk_info = get_members(rtype)
                        
                        comps = []
                        if info:
                            for p in info.get("properties", {}):
                                comps.append({"label": p, "kind": "Property", "detail": f"Property of {rtype}"})
                            for m in info.get("methods", {}):
                                comps.append({"label": m, "kind": "Method", "detail": f"Method of {rtype}"})
                        
                        if sdk_info:
                            for p in sdk_info.get("properties", []):
                                comps.append({"label": p, "kind": "Property", "detail": f"SDK Property: {rtype}.{p}"})
                            for m in sdk_info.get("methods", []):
                                comps.append({"label": m, "kind": "Method", "detail": f"SDK Method: {rtype}.{m}"})
                        
                        return comps
            except Exception:
                pass
        else:
            # Global symbol completion
            import re
            match = re.search(r"([a-zA-Z_0-9]+)$", prefix)
            if not match:
                # If we are at the start of a line or after a space, suggest everything
                word = ""
            else:
                word = match.group(1)

            comps = []
            
            # 1. Add SDK Globals
            from transpiler.infrastructure.sdk_metadata import SDK_GLOBALS
            for g in SDK_GLOBALS:
                if g.startswith(word):
                    kind = "Variable"
                    if g in ("print", "warn", "error", "spawn", "delay", "wait"): kind = "Function"
                    elif g[0].isupper() and g not in ("Workspace",): kind = "Class"
                    comps.append({"label": g, "kind": kind, "detail": "Roblox Global"})

            # 2. Add local symbols from the module being edited
            if cached and cached.get("module"):
                # We can use the cached module symbols (but it might be slightly out of date)
                # For better results, we could re-run an analyzer but only on the symbols
                # Let's use the symbols already identified in reload_file
                # Actually, QueryService doesn't store the symbol table, it stores the ModuleExport
                mod_name = self._get_module_name(Path(uri.replace("file:///", "").replace("file://", "")))
                if self.project_context:
                    mod_export = self.project_context.modules.get(mod_name)
                    if mod_export:
                        for s, t in mod_export.symbols.items():
                            if s.startswith(word):
                                kind = "Variable"
                                if t.startswith("Class:"): kind = "Class"
                                elif "Function" in t or "method" in t: kind = "Function"
                                comps.append({"label": s, "kind": kind, "detail": f"Local {t}"})

            return comps
        
        return []

    def _get_node_at(self, uri: str, line: int, col: int) -> Optional[ir.IRNode]:
        cached = self._cache.get(uri)
        if not cached or not cached["module"]:
            return None
            
        ir_module = cached["module"]
        finder = NodeFinder(line, col)
        finder.visit(ir_module)
        return finder.found_node

    def get_definition(self, uri: str, line: int, col: int) -> Optional[Dict[str, Any]]:
        uri = self._normalize_uri(uri)
        node = self._get_node_at(uri, line, col)
        if not node:
            return None

        # Handle IRFunctionCall by jumping to the function itself
        if isinstance(node, ir.IRFunctionCall) and isinstance(node.func, ir.IRVariable):
            node = node.func
            
        # 1. Handle Variables (including imported ones)
        if isinstance(node, ir.IRVariable):
            defn_info = getattr(node, "definition_location", None)
            if defn_info:
                def_line, def_col, origin = defn_info
                target_uri = uri
                
                # If it's an imported symbol, find its origin in ProjectContext
                if origin and self.project_context:
                    mod_export = self.project_context.modules.get(origin)
                    if mod_export and mod_export.uri:
                        target_uri = mod_export.uri
                        # Use the definition location from the source module's exports
                        if node.name in mod_export.definitions:
                            def_line, def_col = mod_export.definitions[node.name]
                
                return {
                    "uri": target_uri,
                    "range": {
                        "start": {"line": def_line - 1, "character": def_col},
                        "end": {"line": def_line - 1, "character": def_col + len(node.name)}
                    }
                }
        
        # 2. Handle Property Access (jump to class definition)
        elif isinstance(node, ir.IRPropertyAccess):
            receiver_type = node.receiver.inferred_type if hasattr(node.receiver, "inferred_type") else None
            if receiver_type and receiver_type.startswith("Class:"):
                class_name = receiver_type[6:]
                if self.project_context:
                    for mod_name, mod_export in self.project_context.modules.items():
                        if class_name in mod_export.definitions:
                            def_line, def_col = mod_export.definitions[class_name]
                            return {
                                "uri": mod_export.uri,
                                "range": {
                                    "start": {"line": def_line - 1, "character": def_col},
                                    "end": {"line": def_line - 1, "character": def_col + len(class_name)}
                                }
                            }
        
        # 3. Handle Imports (jump to the imported module or symbol)
        elif isinstance(node, ir.IRImport):
            if self.project_context:
                mod_export = self.project_context.modules.get(node.module)
                if mod_export and mod_export.uri:
                    return {
                        "uri": mod_export.uri,
                        "range": {"start": {"line": 0, "character": 0}, "end": {"line": 0, "character": 0}}
                    }
        
        elif isinstance(node, ir.IRImportFrom):
            if self.project_context:
                mod_export = self.project_context.modules.get(node.module)
                if mod_export and mod_export.uri:
                    target_uri = mod_export.uri
                    # We could refine this to jump to specific names in the import line
                    return {
                        "uri": target_uri,
                        "range": {"start": {"line": 0, "character": 0}, "end": {"line": 0, "character": 0}}
                    }
        
        return None

    def find_references(self, uri: str, line: int, col: int) -> List[Dict[str, Any]]:
        uri = self._normalize_uri(uri)
        node = self._get_node_at(uri, line, col)
        if not node or not hasattr(node, "name"): return []
        
        symbol_name = node.name
        print(f"DEBUG: Finding references for '{symbol_name}'...")
        refs = []
        if self.project_context:
            for mod_name, mod in self.project_context.modules.items():
                if symbol_name in mod.usages:
                    print(f"DEBUG: Found {len(mod.usages[symbol_name])} usages in module '{mod_name}'")
                    for u, l, c in mod.usages[symbol_name]:
                        refs.append({
                            "uri": u,
                            "range": {
                                "start": {"line": l - 1, "character": c},
                                "end": {"line": l - 1, "character": c + len(symbol_name)}
                            }
                        })
        return refs

query_service = QueryService()
