from __future__ import annotations

import asyncio
import json
import logging
import subprocess
from pathlib import Path
from typing import Any, Callable, Protocol

from pydantic_ai import DeferredToolRequests, DeferredToolResults, ToolDenied

from invisio_graph.git_utils import (
    cleanup_directory,
    clone_repository,
    create_temp_directory,
    is_valid_git_url,
)


logger = logging.getLogger(__name__)
__all__ = ["CodeGraph", "RagSession"]


# Optional module-level aliases kept for compatibility with monkeypatching/tests.
connect_memgraph: Callable[[int], Any] | None = None
load_parsers: Callable[[], Any] | None = None
GraphUpdater: type[Any] | None = None
CypherGenerator: type[Any] | None = None
CodeRetriever: type[Any] | None = None
embed_code: Callable[..., Any] | None = None
search_embeddings: Callable[..., Any] | None = None
initialize_services_and_agent: Callable[[str, Any], Any] | None = None


def _get_settings() -> Any:
    """Fetches the settings like model config cricignore and other stuff"""
    from codebase_rag.config import settings

    return settings


def _get_default_connect_fn() -> Callable[[int], Any]:
    """Connects to a memgraph instance"""
    global connect_memgraph
    if connect_memgraph is None:
        from codebase_rag.main import connect_memgraph as _connect_memgraph

        connect_memgraph = _connect_memgraph
    return connect_memgraph


def _get_default_parser_loader() -> Callable[[], Any]:
    """Load Parsers"""
    global load_parsers
    if load_parsers is None:
        from codebase_rag.parser_loader import load_parsers as _load_parsers

        load_parsers = _load_parsers
    return load_parsers


def _get_default_graph_updater_cls() -> type[Any]:
    """Instainiate the GraphUpdater class that is used for graph operations"""
    global GraphUpdater
    if GraphUpdater is None:
        from codebase_rag.graph_updater import GraphUpdater as _graph_updater_cls

        GraphUpdater = _graph_updater_cls
    return GraphUpdater


def _get_cypher_generator_cls() -> type[Any]:
    """Get the nlp to cypher model"""
    global CypherGenerator
    if CypherGenerator is None:
        from codebase_rag.services.llm import CypherGenerator as _cypher_generator_cls

        CypherGenerator = _cypher_generator_cls
    return CypherGenerator


def _get_code_retriever_cls() -> type[Any]:
    global CodeRetriever
    if CodeRetriever is None:
        from codebase_rag.tools.code_retrieval import CodeRetriever as _code_retriever_cls

        CodeRetriever = _code_retriever_cls
    return CodeRetriever


def _get_embed_code_fn() -> Callable[..., Any]:
    global embed_code
    if embed_code is None:
        from codebase_rag.embedder import embed_code as _embed_code

        embed_code = _embed_code
    return embed_code


def _get_search_embeddings_fn() -> Callable[..., Any]:
    global search_embeddings
    if search_embeddings is None:
        from codebase_rag.vector_store import search_embeddings as _search_embeddings

        search_embeddings = _search_embeddings
    return search_embeddings


def _build_nodes_by_ids_query(node_ids: list[int]) -> str:
    from codebase_rag.cypher_queries import build_nodes_by_ids_query

    return build_nodes_by_ids_query(node_ids)


def _get_agent_initializer() -> Callable[[str, Any], Any]:
    global initialize_services_and_agent
    if initialize_services_and_agent is None:
        from codebase_rag.main import (
            _initialize_services_and_agent as _initialize_services_and_agent,
        )

        initialize_services_and_agent = _initialize_services_and_agent
    return initialize_services_and_agent


class _IngestorProtocol(Protocol):
    def clean_database(self) -> None: ...

    def ensure_constraints(self) -> None: ...

    def fetch_all(self, query: str, params: dict[str, Any] | None = None) -> list: ...

    def list_projects(self) -> list[str]: ...

    def delete_project(self, project_name: str) -> None: ...


class RagSession:
    """Headless RAG chat session for UI clients (CLI, web, React, etc.)."""

    def __init__(
        self,
        *,
        project_root: Path,
        context_manager: Any,
        rag_agent: Any,
    ) -> None:
        self.project_root = project_root
        self._context_manager = context_manager
        self._rag_agent = rag_agent
        self._message_history: list[Any] = []
        self._pending_requests: DeferredToolRequests | None = None
        self._active_question: str | None = None
        self._closed = False

    def _require_open(self) -> None:
        if self._closed:
            raise RuntimeError("RAG session is closed")

    @staticmethod
    def _require_non_empty_question(question: str) -> str:
        if not isinstance(question, str):
            raise TypeError("question must be a string")
        cleaned = question.strip()
        if not cleaned:
            raise ValueError("question cannot be empty")
        return cleaned

    @staticmethod
    def _serialize_tool_requests(requests: DeferredToolRequests) -> list[dict[str, Any]]:
        pending: list[dict[str, Any]] = []
        for call in requests.approvals:
            pending.append(
                {
                    "tool_call_id": call.tool_call_id,
                    "tool_name": call.tool_name,
                    "args": call.args_as_dict(),
                }
            )
        return pending

    def _build_deferred_results(
        self, decisions: dict[str, bool | str]
    ) -> DeferredToolResults:
        if not isinstance(decisions, dict):
            raise TypeError("decisions must be a dictionary")
        if self._pending_requests is None:
            raise ValueError("No pending tool approvals")

        deferred_results = DeferredToolResults()
        for call in self._pending_requests.approvals:
            if call.tool_call_id not in decisions:
                raise ValueError(f"Missing decision for tool call: {call.tool_call_id}")
            decision = decisions[call.tool_call_id]
            if isinstance(decision, bool):
                deferred_results.approvals[call.tool_call_id] = decision
            elif isinstance(decision, str):
                message = decision.strip() or "Tool request denied by client"
                deferred_results.approvals[call.tool_call_id] = ToolDenied(message)
            else:
                raise TypeError(
                    "Each decision value must be bool (approve/deny) or str (deny reason)"
                )
        return deferred_results

    async def ask(self, question: str) -> dict[str, Any]:
        """Send a new user message and get either assistant output or tool approvals."""
        self._require_open()
        if self._pending_requests is not None:
            raise ValueError("Pending tool approvals must be resolved before asking again")

        cleaned_question = self._require_non_empty_question(question)
        self._active_question = cleaned_question

        response = await self._rag_agent.run(
            cleaned_question,
            message_history=self._message_history,
            deferred_tool_results=None,
        )
        self._message_history.extend(response.new_messages())

        if isinstance(response.output, DeferredToolRequests):
            self._pending_requests = response.output
            return {
                "status": "approval_required",
                "question": cleaned_question,
                "pending_approvals": self._serialize_tool_requests(response.output),
            }

        output_text = response.output if isinstance(response.output, str) else ""
        self._pending_requests = None
        self._active_question = None
        return {
            "status": "completed",
            "question": cleaned_question,
            "assistant_message": output_text,
        }

    async def resolve_tool_approvals(
        self, decisions: dict[str, bool | str]
    ) -> dict[str, Any]:
        """Continue a pending turn by approving/denying requested tool calls."""
        self._require_open()
        if self._pending_requests is None or self._active_question is None:
            raise ValueError("No pending tool approvals")

        deferred_results = self._build_deferred_results(decisions)
        response = await self._rag_agent.run(
            self._active_question,
            message_history=self._message_history,
            deferred_tool_results=deferred_results,
        )
        self._message_history.extend(response.new_messages())

        if isinstance(response.output, DeferredToolRequests):
            self._pending_requests = response.output
            return {
                "status": "approval_required",
                "question": self._active_question,
                "pending_approvals": self._serialize_tool_requests(response.output),
            }

        output_text = response.output if isinstance(response.output, str) else ""
        finished_question = self._active_question
        self._pending_requests = None
        self._active_question = None
        return {
            "status": "completed",
            "question": finished_question,
            "assistant_message": output_text,
        }

    def close(self) -> None:
        if not self._closed:
            self._context_manager.__exit__(None, None, None)
            self._closed = True

    def __enter__(self) -> RagSession:
        return self

    def __exit__(self, exc_type, exc, tb) -> bool:
        self.close()
        return False


class CodeGraph:
    """Public SDK client for indexing and querying code graphs.

    Public API stability:
    - Methods without a leading underscore are considered part of the stable SDK contract.
    - Methods prefixed with "_" are internal helpers and may change without notice.
    """

    def __init__(
        self,
        batch_size: int | None = None,
        *,
        connect_fn: Callable[[int], Any] | None = None,
        parser_loader: Callable[[], Any] | None = None,
        graph_updater_cls: type[Any] | None = None,
    ) -> None:
        settings = _get_settings()
        resolved = settings.resolve_batch_size(batch_size)
        self._batch_size = resolved if resolved else 200
        self._connect_fn = connect_fn or _get_default_connect_fn()
        self._parser_loader = parser_loader or _get_default_parser_loader()
        self._graph_updater_cls = graph_updater_cls or _get_default_graph_updater_cls()

    def _require_non_empty_text(self, value: str, name: str) -> str:
        """Basic input sanitization"""
        if not isinstance(value, str):
            raise TypeError(f"{name} must be a string")
        text = value.strip()
        if not text:
            raise ValueError(f"{name} cannot be empty")
        return text

    def _resolve_existing_directory(self, repo_path: str) -> Path:
        """Validates directory path"""
        normalized = self._require_non_empty_text(repo_path, "repo_path")
        path = Path(normalized).expanduser().resolve()
        if not path.exists():
            raise FileNotFoundError(f"Repository path does not exist: {path}")
        if not path.is_dir():
            raise NotADirectoryError(f"Repository path is not a directory: {path}")
        return path

    def _resolve_project_and_target(
        self, project_root: str, relative_path: str, *, action: str
    ) -> tuple[Path, Path]:
        project = Path(self._require_non_empty_text(project_root, "project_root"))
        project = project.expanduser().resolve()
        if not project.exists():
            raise FileNotFoundError(f"Project root does not exist: {project}")
        if not project.is_dir():
            raise NotADirectoryError(f"Project root is not a directory: {project}")

        target_input = self._require_non_empty_text(relative_path, action)
        target = (project / target_input).resolve()
        try:
            target.relative_to(project)
        except ValueError as exc:
            raise PermissionError(
                f"Path escapes project root for {action}: {relative_path!r}"
            ) from exc
        return project, target

    def _normalize_rows(self, rows: list[Any], *, context: str) -> list[dict[str, Any]]:
        normalized: list[dict[str, Any]] = []
        for index, row in enumerate(rows):
            if isinstance(row, dict):
                normalized.append(row)
                continue
            try:
                normalized.append(dict(row))
            except Exception as exc:
                raise ValueError(
                    f"Received malformed graph data in {context} at row index {index}"
                ) from exc
        return normalized

    # =====================
    # INDEXING
    # =====================

    def update(self, repo_path: str, clean: bool = False) -> None:
        """Index or re-index a local repository into the graph database."""
        resolved_repo_path = self._resolve_existing_directory(repo_path)
        logger.debug(
            "Starting graph update",
            extra={"repo_path": str(resolved_repo_path), "clean": clean},
        )
        with self._connect_fn(self._batch_size) as ingestor:
            if clean:
                ingestor.clean_database()

            ingestor.ensure_constraints()
            parsed = self._parser_loader()
            if not isinstance(parsed, tuple) or len(parsed) != 2:
                raise RuntimeError("Parser loader returned invalid parser metadata")
            parsers, queries = parsed

            updater = self._graph_updater_cls(
                ingestor=ingestor,
                repo_path=resolved_repo_path,
                parsers=parsers,
                queries=queries,
            )
            try:
                updater.run()
            except Exception as exc:
                raise RuntimeError(
                    f"Failed to update graph for repository: {resolved_repo_path}"
                ) from exc
            logger.debug("Graph update completed", extra={"repo_path": str(resolved_repo_path)})

    def update_from_git(self, repo_url: str, clean: bool = False) -> None:
        """Clone a git repository and index it."""
        repo_url = self._require_non_empty_text(repo_url, "repo_url")
        if not is_valid_git_url(repo_url):
            raise ValueError(f"Invalid git URL: {repo_url}")
        repo_name = repo_url.rstrip("/").split("/")[-1]
        if repo_name.endswith(".git"):
            repo_name = repo_name[:-4]
        temp_dir = create_temp_directory()
        repo_path = temp_dir / repo_name
        logger.debug(
            "Cloning repository before update",
            extra={"repo_url": repo_url, "temp_repo_path": str(repo_path)},
        )

        try:
            if not clone_repository(repo_url, repo_path):
                raise RuntimeError("Failed to clone repository")

            self.update(str(repo_path), clean=clean)
        finally:
            cleanup_directory(temp_dir)
            logger.debug("Temporary repository directory cleaned", extra={"path": str(temp_dir)})

    # =====================
    # GRAPH QUERYING
    # =====================

    async def query_graph(self, question: str) -> list[dict[str, Any]]:
        """Run natural-language graph queries and return normalized row dictionaries."""
        question = self._require_non_empty_text(question, "question")
        with self._connect_fn(self._batch_size) as ingestor:
            cypher_generator_cls = _get_cypher_generator_cls()
            generator = cypher_generator_cls()
            cypher_query = await generator.generate(question)
            results = ingestor.fetch_all(cypher_query)
            return self._normalize_rows(results, context="query_graph")

    def start_rag_session(self, repo_path: str) -> RagSession:
        """Create a headless RAG session suitable for web/React clients."""
        resolved_repo_path = self._resolve_existing_directory(repo_path)
        context_manager = self._connect_fn(self._batch_size)
        ingestor = context_manager.__enter__()
        initializer = _get_agent_initializer()
        try:
            rag_agent, _tool_names = initializer(str(resolved_repo_path), ingestor)
        except Exception:
            context_manager.__exit__(None, None, None)
            raise

        return RagSession(
            project_root=resolved_repo_path,
            context_manager=context_manager,
            rag_agent=rag_agent,
        )

    def cypher(
        self, query: str, params: dict[str, Any] | None = None
    ) -> list[dict[str, Any]]:
        """Execute a raw Cypher query and return normalized row dictionaries."""
        query = self._require_non_empty_text(query, "query")
        if params is not None and not isinstance(params, dict):
            raise TypeError("params must be a dictionary when provided")
        with self._connect_fn(self._batch_size) as ingestor:
            results = ingestor.fetch_all(query, params or {})
            return self._normalize_rows(results, context="cypher")

    def export_graph_json(self, output_path: str) -> dict[str, Any]:
        """Export graph data from Memgraph to a JSON file."""
        output_path = self._require_non_empty_text(output_path, "output_path")
        target = Path(output_path).expanduser().resolve()

        with self._connect_fn(self._batch_size) as ingestor:
            export_fn = getattr(ingestor, "export_graph_to_dict", None)
            if not callable(export_fn):
                raise RuntimeError("Connected ingestor does not support graph export")
            graph_data = export_fn()

        target.parent.mkdir(parents=True, exist_ok=True)
        with target.open("w", encoding="utf-8") as handle:
            json.dump(graph_data, handle, indent=2, ensure_ascii=False)

        return {
            "graph":graph_data,
            "output_path": str(target),
            "total_nodes": graph_data.get("metadata", {}).get("total_nodes", 0),
            "total_relationships": graph_data.get("metadata", {}).get(
                "total_relationships", 0
            ),
        }

    def export_project(self, project_name: str) -> dict[str, Any]:
        """Return a single project's graph data from Memgraph as a dictionary."""
        project_name = self._require_non_empty_text(project_name, "project_name")

        with self._connect_fn(self._batch_size) as ingestor:
            export_fn = getattr(ingestor, "export_project_to_dict", None)
            if not callable(export_fn):
                raise RuntimeError("Connected ingestor does not support project export")
            graph_data = export_fn(project_name)

        return {
            "project_name": project_name,
            "graph": graph_data,
            "total_nodes": graph_data.get("metadata", {}).get("total_nodes", 0),
            "total_relationships": graph_data.get("metadata", {}).get(
                "total_relationships", 0
            ),
        }

    def run_rag_terminal(self, repo_path: str) -> None:
        """Run the interactive terminal RAG session (equivalent to `cgr start`)."""
        resolved_repo_path = self._resolve_existing_directory(repo_path)
        from codebase_rag.main import main_async

        try:
            asyncio.run(main_async(str(resolved_repo_path), self._batch_size))
        except RuntimeError as exc:
            if "asyncio.run() cannot be called from a running event loop" not in str(exc):
                raise
            raise RuntimeError(
                "run_rag_terminal must be called from a synchronous context"
            ) from exc

    # =====================
    # PROJECT MANAGEMENT
    # =====================

    def list_projects(self) -> list[str]:
        """List indexed project names."""
        with self._connect_fn(self._batch_size) as ingestor:
            return ingestor.list_projects()

    def delete_project(self, project_name: str) -> None:
        """Delete an indexed project by name."""
        project_name = self._require_non_empty_text(project_name, "project_name")
        with self._connect_fn(self._batch_size) as ingestor:
            ingestor.delete_project(project_name)

    def wipe_database(self) -> None:
        """Delete all indexed data from the backing graph database."""
        with self._connect_fn(self._batch_size) as ingestor:
            ingestor.clean_database()

    # =====================
    # CODE RETRIEVAL
    # =====================

    async def get_code_snippet(
        self, qualified_name: str, project_root: str
    ) -> dict[str, Any]:
        """Get code for a qualified symbol as a plain dictionary response."""
        qualified_name = self._require_non_empty_text(qualified_name, "qualified_name")
        self._resolve_existing_directory(project_root)
        with self._connect_fn(self._batch_size) as ingestor:
            code_retriever_cls = _get_code_retriever_cls()
            retriever = code_retriever_cls(
                project_root=project_root,
                ingestor=ingestor,
            )
            snippet = await retriever.find_code_snippet(qualified_name)

        if isinstance(snippet, dict):
            return dict(snippet)
        if hasattr(snippet, "model_dump"):
            return dict(snippet.model_dump())
        if hasattr(snippet, "dict"):
            return dict(snippet.dict())
        return {
            "qualified_name": qualified_name,
            "source_code": getattr(snippet, "source_code", ""),
            "file_path": getattr(snippet, "file_path", ""),
            "line_start": getattr(snippet, "line_start", 0),
            "line_end": getattr(snippet, "line_end", 0),
            "docstring": getattr(snippet, "docstring", None),
            "found": getattr(snippet, "found", False),
            "error_message": getattr(snippet, "error_message", None),
        }

    # =====================
    # SEMANTIC SEARCH
    # =====================

    def semantic_search(self, query: str, top_k: int = 5) -> list[dict[str, Any]]:
        """Run semantic search and return ranked matches as plain dictionaries."""
        query = self._require_non_empty_text(query, "query")
        if not isinstance(top_k, int):
            raise TypeError("top_k must be an integer")
        if top_k < 1:
            raise ValueError("top_k must be >= 1")
        embed_code_fn = _get_embed_code_fn()
        search_embeddings_fn = _get_search_embeddings_fn()
        embedding = embed_code_fn(query)
        raw_results = search_embeddings_fn(embedding, top_k)

        normalized_hits: list[dict[str, Any]] = []
        for item in raw_results or []:
            if isinstance(item, dict):
                normalized_hits.append(dict(item))
                continue
            if isinstance(item, tuple) and len(item) == 2:
                normalized_hits.append({"node_id": item[0], "score": item[1]})
                continue
            normalized_hits.append({"value": item})

        node_ids = [
            item["node_id"]
            for item in normalized_hits
            if isinstance(item.get("node_id"), int)
        ]
        if not node_ids:
            return normalized_hits

        query_text = _build_nodes_by_ids_query(node_ids)
        params = {str(i): node_id for i, node_id in enumerate(node_ids)}

        with self._connect_fn(self._batch_size) as ingestor:
            rows = self._normalize_rows(
                ingestor.fetch_all(query_text, params),
                context="semantic_search",
            )

        metadata_by_id = {
            row["node_id"]: row
            for row in rows
            if isinstance(row.get("node_id"), int)
        }

        enriched_results: list[dict[str, Any]] = []
        for hit in normalized_hits:
            node_id = hit.get("node_id")
            if not isinstance(node_id, int) or node_id not in metadata_by_id:
                enriched_results.append(hit)
                continue

            metadata = metadata_by_id[node_id]
            enriched = dict(hit)
            enriched.update(
                {
                    "qualified_name": metadata.get("qualified_name"),
                    "name": metadata.get("name"),
                    "type": metadata.get("type"),
                }
            )
            enriched_results.append(enriched)

        return enriched_results

    # =====================
    # FILE OPERATIONS
    # =====================

    def read_file(
        self,
        file_path: str,
        project_root: str,
        offset: int | None = None,
        limit: int | None = None,
    ) -> str:
        """Read a UTF-8 file under project_root with optional line pagination."""
        _, target = self._resolve_project_and_target(
            project_root, file_path, action="file_path"
        )
        if not target.exists():
            raise FileNotFoundError(f"File does not exist: {target}")
        if not target.is_file():
            raise IsADirectoryError(f"Path is not a file: {target}")

        if offset is not None and (not isinstance(offset, int) or offset < 0):
            raise ValueError("offset must be a non-negative integer when provided")
        if limit is not None and (not isinstance(limit, int) or limit < 0):
            raise ValueError("limit must be a non-negative integer when provided")

        content = target.read_text(encoding="utf-8")
        lines = content.splitlines()

        if offset is not None:
            lines = lines[offset:]
        if limit is not None:
            lines = lines[:limit]

        return "\n".join(lines)

    def write_file(self, file_path: str, content: str, project_root: str) -> None:
        """Write a UTF-8 file under project_root, creating parent directories as needed."""
        _, target = self._resolve_project_and_target(
            project_root, file_path, action="file_path"
        )
        if not isinstance(content, str):
            raise TypeError("content must be a string")

        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(content, encoding="utf-8")

    def replace_code(
        self,
        file_path: str,
        target_code: str,
        replacement_code: str,
        project_root: str,
    ) -> None:
        """Replace the first matching code fragment in a UTF-8 file under project_root."""
        _, target = self._resolve_project_and_target(
            project_root, file_path, action="file_path"
        )
        if not target.exists():
            raise FileNotFoundError(f"File does not exist: {target}")
        if not target.is_file():
            raise IsADirectoryError(f"Path is not a file: {target}")
        if not isinstance(target_code, str) or not target_code:
            raise ValueError("target_code must be a non-empty string")
        if not isinstance(replacement_code, str):
            raise TypeError("replacement_code must be a string")

        content = target.read_text(encoding="utf-8")
        if target_code not in content:
            raise ValueError("target_code was not found in file content")
        new_content = content.replace(target_code, replacement_code, 1)
        target.write_text(new_content, encoding="utf-8")

    # =====================
    # DIRECTORY
    # =====================

    def list_directory(
        self, directory_path: str, project_root: str
    ) -> list[dict[str, str]]:
        """List direct children of a directory under project_root."""
        _, target = self._resolve_project_and_target(
            project_root, directory_path, action="directory_path"
        )
        if not target.exists():
            raise FileNotFoundError(f"Directory does not exist: {target}")
        if not target.is_dir():
            raise NotADirectoryError(f"Path is not a directory: {target}")

        entries = []
        for item in sorted(target.iterdir(), key=lambda entry: entry.name.lower()):
            rel_path = item.relative_to(target)
            entries.append(
                {
                    "name": item.name,
                    "path": str(rel_path).replace("\\", "/"),
                    "type": "directory" if item.is_dir() else "file",
                }
            )

        return entries

    # =====================
    # SHELL
    # =====================

    def execute_shell(
        self,
        command: str,
        project_root: str,
        timeout: int = 30,
    ) -> dict[str, Any]:
        """Execute an allowlisted shell command scoped to project_root."""
        command = self._require_non_empty_text(command, "command")
        if not isinstance(timeout, int) or timeout < 1:
            raise ValueError("timeout must be a positive integer")
        project = self._resolve_existing_directory(project_root)
        settings = _get_settings()
        allowed = settings.SHELL_COMMAND_ALLOWLIST
        cmd_parts = command.strip().split()

        if not cmd_parts or cmd_parts[0] not in allowed:
            raise ValueError("Command not allowed")
        if any(token in {"&&", "||", ";", "|", ">", ">>", "<"} for token in cmd_parts):
            raise ValueError("Shell control operators are not allowed")

        result = subprocess.run(
            command,
            shell=True,
            cwd=project,
            capture_output=True,
            text=True,
            timeout=timeout,
        )

        return {
            "return_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr,
        }
