"""Rosetta MCP V2 server assembly."""

from __future__ import annotations

import base64
from importlib import resources as pkg_resources
from typing import Any

from mcp.server.fastmcp import Context, FastMCP
from mcp.types import Icon

from ims_mcp.analytics.tracker import debug_print, register_signal_handlers, set_runtime_config, track_tool_call
from ims_mcp.analytics.user_context import get_repository_from_context, get_username
from ims_mcp.clients.dataset import DatasetLookup
from ims_mcp.clients.document import DocumentClient
from ims_mcp.clients.ragflow import RagflowClient
from ims_mcp.config import RosettaConfig
from ims_mcp.constants import (
    TAG_STARTUP,
    TOOL_DISCOVER_PROJECTS,
    TOOL_GET_CONTEXT_INSTRUCTIONS,
    TOOL_QUERY_INSTRUCTIONS,
    TOOL_QUERY_PROJECT_CONTEXT,
    TOOL_STORE_PROJECT_CONTEXT,
    TOOL_SUBMIT_FEEDBACK,
)
from ims_mcp.context import CallContext
from ims_mcp.services.bundler import Bundler
from ims_mcp.services.feedback import FeedbackService
from ims_mcp.services.query_builder import QueryBuilder
from ims_mcp.tool_prompts import (
    PROMPT_DISCOVER_PROJECTS,
    PROMPT_GET_CONTEXT_INSTRUCTIONS,
    PROMPT_QUERY_INSTRUCTIONS,
    PROMPT_QUERY_PROJECT_CONTEXT,
    PROMPT_STORE_PROJECT_CONTEXT,
    PROMPT_SUBMIT_FEEDBACK,
)
from ims_mcp.tools.feedback import submit_feedback as submit_feedback_tool
from ims_mcp.tools.instructions import (
    get_context_instructions as get_context_instructions_tool,
)
from ims_mcp.tools.instructions import query_instructions as query_instructions_tool
from ims_mcp.tools.projects import (
    discover_projects as discover_projects_tool,
)
from ims_mcp.tools.projects import (
    query_project_context as query_project_context_tool,
)
from ims_mcp.tools.projects import (
    store_project_context as store_project_context_tool,
)

_CONFIG = RosettaConfig.from_env()
set_runtime_config(_CONFIG)
register_signal_handlers()

_RAGFLOW = RagflowClient(_CONFIG).client if _CONFIG.api_key else None
_DATASET_LOOKUP = DatasetLookup(_RAGFLOW) if _RAGFLOW else None
_DOCUMENT_CLIENT = DocumentClient()
_QUERY_BUILDER = QueryBuilder()
_BUNDLER = Bundler(_DOCUMENT_CLIENT)
_FEEDBACK = FeedbackService()


def load_rosetta_icon() -> Icon | None:
    try:
        ref = pkg_resources.files("ims_mcp.resources").joinpath("rosetta-icon.png")
        icon_data = ref.read_bytes()
        data_uri = f"data:image/png;base64,{base64.b64encode(icon_data).decode('utf-8')}"
        return Icon(src=data_uri, mimeType="image/png", sizes=["160x160"])
    except Exception:
        return None


def _load_startup_instructions() -> str:
    if not _RAGFLOW or not _DATASET_LOOKUP:
        return ""
    try:
        dataset_name = _CONFIG.instruction_dataset
        dataset_id = _DATASET_LOOKUP.get_id(dataset_name)
        if not dataset_id:
            return ""
        dataset = _RAGFLOW.get_dataset(name=dataset_name)
        docs = _DOCUMENT_CLIENT.list_docs(
            dataset=dataset,
            page_size=1000,
            **_QUERY_BUILDER.build_list_params(tags=[TAG_STARTUP]),
        )
        return _BUNDLER.bundle(docs, dataset_name) if docs else ""
    except Exception as exc:
        debug_print(_CONFIG, f"[ims-mcp] startup instructions load failed: {exc}")
        return ""


_STARTUP_CONTENT = _load_startup_instructions()
_ROSETTA_ICON = load_rosetta_icon()

mcp = FastMCP(
    name="Rosetta",
    instructions=_STARTUP_CONTENT,
    icons=[_ROSETTA_ICON] if _ROSETTA_ICON else [],
)


async def _log(ctx: Context | None, level: str, message: str) -> None:
    if not ctx:
        return
    try:
        fn = getattr(ctx, level)
        await fn(message)
    except Exception:
        pass


async def _build_call_context(tool_name: str, params: dict[str, Any], ctx: Context) -> CallContext:
    return CallContext(
        config=_CONFIG,
        ragflow=_RAGFLOW,
        dataset_lookup=_DATASET_LOOKUP,
        ctx=ctx,
        username=get_username(),
        repository=await get_repository_from_context(ctx),
        tool_name=tool_name,
        params=params,
    )


def _validate_topic(topic: str | None) -> str | None:
    if topic and len(topic.split()) > 10:
        return "Error: topic must be 10 words or less"
    return None


def _normalize_tags(tags: list[str] | str | None) -> list[str] | None:
    """Convert single string to list, keep list as-is, None stays None."""
    if tags is None:
        return None
    if isinstance(tags, str):
        return [tags] if tags.strip() else None
    # Already a list
    return tags if tags else None


async def _retry_once(fn):
    last_exc = None
    for _ in range(2):
        try:
            return await fn()
        except Exception as exc:
            last_exc = exc
            if _RAGFLOW is None:
                break
    raise RuntimeError(str(last_exc))


@mcp.tool(name=TOOL_GET_CONTEXT_INSTRUCTIONS, description=PROMPT_GET_CONTEXT_INSTRUCTIONS)
@track_tool_call
async def get_context_instructions(topic: str | None = None, ctx: Context | None = None) -> str:
    """
    Args:
        topic: Optional intent summary (max 10 words) for semantic search (slow).

    Returns:
        Bundled instruction documents containing context, rules, and task-specific guidance.
        Includes bootstrap instructions, agent rules, and request classification.
    """
    
    if not _RAGFLOW:
        return "Error: ROSETTA_API_KEY is required"
    err = _validate_topic(topic)
    if err:
        return err

    await _log(ctx, "info", "Loading context instructions")
    call_ctx = await _build_call_context(TOOL_GET_CONTEXT_INSTRUCTIONS, {"topic": topic}, ctx)
    return await _retry_once(
        lambda: get_context_instructions_tool(
            call_ctx=call_ctx,
            document_client=_DOCUMENT_CLIENT,
            bundler=_BUNDLER,
            query_builder=_QUERY_BUILDER,
            topic=topic,
        )
    )


@mcp.tool(name=TOOL_QUERY_INSTRUCTIONS, description=PROMPT_QUERY_INSTRUCTIONS)
@track_tool_call
async def query_instructions(
    query: str | None = None,
    tags: list[str] | str | None = None,
    topic: str | None = None,
    ctx: Context | None = None,
) -> str:
    """
    Args:
        query: Text to search for in instruction documents.
        tags: Filter by instruction tags (string or list). Prefer tags when instruction family is known.
        topic: Optional intent summary (max 10 words) for semantic search (slow).

    Returns:
        Matching instruction documents with metadata, content chunks, and context.
        Returns empty result if no documents match the criteria.
    """
    if not _RAGFLOW:
        return "Error: ROSETTA_API_KEY is required"
    err = _validate_topic(topic)
    if err:
        return err

    normalized_tags = _normalize_tags(tags)
    await _log(ctx, "info", "Querying instructions")
    call_ctx = await _build_call_context(
        TOOL_QUERY_INSTRUCTIONS,
        {"query": query, "tags": normalized_tags, "topic": topic},
        ctx,
    )
    return await _retry_once(
        lambda: query_instructions_tool(
            call_ctx=call_ctx,
            document_client=_DOCUMENT_CLIENT,
            bundler=_BUNDLER,
            query_builder=_QUERY_BUILDER,
            query=query,
            tags=normalized_tags,
            topic=topic,
        )
    )


@mcp.tool(name=TOOL_SUBMIT_FEEDBACK, description=PROMPT_SUBMIT_FEEDBACK)
@track_tool_call
async def submit_feedback(request_mode: str, feedback: dict, ctx: Context | None = None) -> str:
    """
    Args:
        request_mode: The AI agent mode/classification used for this request. Examples: "coding.md", "help.md", "research.md", "aqa.md".
        feedback: Structured brief feedback dictionary with keys:
                  - summary (required): Brief description of the feedback (1-2 sentences)
                  - root_cause (required): Why the issue occurred or what needs improvement
                  - prompt_suggestions (required): Specific suggestions for improvement of instructions
                  - context (required): Context where this feedback applies
                  - Additional keys are allowed for extended feedback

    Returns:
        Confirmation of feedback submission and any actionable insights.
    """

    if not _RAGFLOW:
        return "Error: ROSETTA_API_KEY is required"
    await _log(ctx, "info", "Submitting feedback")
    call_ctx = await _build_call_context(
        TOOL_SUBMIT_FEEDBACK,
        {"request_mode": request_mode, "feedback": feedback},
        ctx,
    )
    return await _retry_once(
        lambda: submit_feedback_tool(
            call_ctx=call_ctx,
            feedback_service=_FEEDBACK,
            request_mode=request_mode,
            feedback=feedback,
        )
    )


@mcp.tool(name=TOOL_QUERY_PROJECT_CONTEXT, description=PROMPT_QUERY_PROJECT_CONTEXT)
@track_tool_call
async def query_project_context(
    repository_name: str,
    query: str | None = None,
    tags: list[str] | str | None = None,
    topic: str | None = None,
    ctx: Context | None = None,
) -> str:
    """
    Args:
        repository_name: Name of the workspace folder (or the repository) to query.
        query: Text to search for in project context documents.
        tags: Filter by context tags (string or list).
        topic: Optional intent summary (max 10 words) for semantic search (slow).

    Returns:
        Project context documents including architecture, implementation details, and decisions. Returns empty result if project dataset not found.
    """

    if not _RAGFLOW:
        return "Error: ROSETTA_API_KEY is required"
    err = _validate_topic(topic)
    if err:
        return err

    normalized_tags = _normalize_tags(tags)
    await _log(ctx, "info", f"Querying project context for {repository_name}")
    call_ctx = await _build_call_context(
        TOOL_QUERY_PROJECT_CONTEXT,
        {"repository_name": repository_name, "query": query, "tags": normalized_tags, "topic": topic},
        ctx,
    )
    return await _retry_once(
        lambda: query_project_context_tool(
            call_ctx=call_ctx,
            document_client=_DOCUMENT_CLIENT,
            bundler=_BUNDLER,
            query_builder=_QUERY_BUILDER,
            repository_name=repository_name,
            query=query,
            tags=normalized_tags,
            topic=topic,
        )
    )


@mcp.tool(name=TOOL_STORE_PROJECT_CONTEXT, description=PROMPT_STORE_PROJECT_CONTEXT)
@track_tool_call
async def store_project_context(
    repository_name: str,
    document: str,
    tags: list[str] | str,
    content: str,
    force: bool = False,
    ctx: Context | None = None,
) -> str:
    """
    Args:
        repository_name: Name of the workspace folder (or the repository) to query.
        document: Document name. Examples: "ARCHITECTURE.md"
        tags: Tags to categorize the document. Examples: ["implementation", "api-design"]
        content: The actual content of the document.
        force: Do not force. Try to discover the repository first. If true, create repository dataset if it doesn't exist.

    Returns:
        Confirmation of document storage.
    """

    if not _RAGFLOW:
        return "Error: ROSETTA_API_KEY is required"
    normalized_tags = _normalize_tags(tags) or []
    await _log(ctx, "info", f"Storing project context for {repository_name}")
    call_ctx = await _build_call_context(
        TOOL_STORE_PROJECT_CONTEXT,
        {
            "repository_name": repository_name,
            "document": document,
            "tags": normalized_tags,
            "force": force,
        },
        ctx,
    )
    return await _retry_once(
        lambda: store_project_context_tool(
            call_ctx=call_ctx,
            document_client=_DOCUMENT_CLIENT,
            repository_name=repository_name,
            document=document,
            tags=normalized_tags,
            content=content,
            force=force,
        )
    )


@mcp.tool(name=TOOL_DISCOVER_PROJECTS, description=PROMPT_DISCOVER_PROJECTS)
@track_tool_call
async def discover_projects(query: str | None = None, ctx: Context | None = None) -> str:
    """
    Args:
        query: Optional search term to filter projects by name.

    Returns:
        List of available project datasets with names and metadata.
        Dataset name equals to repository name.
    """

    if not _RAGFLOW:
        return "Error: ROSETTA_API_KEY is required"
    await _log(ctx, "info", "Discovering project datasets")
    call_ctx = await _build_call_context(TOOL_DISCOVER_PROJECTS, {"query": query}, ctx)
    return await _retry_once(lambda: discover_projects_tool(call_ctx=call_ctx, query=query))


def main() -> None:
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
