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

from __future__ import annotations

import base64
import os
import sys
from importlib import resources as pkg_resources
from typing import Annotated, Any

from fastmcp import Context, FastMCP
from mcp.types import Icon
from pydantic import Field

from ims_mcp import __version__ as _MCP_VERSION
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.auth import build_oauth_provider
from ims_mcp.clients.dataset import DatasetLookup
from ims_mcp.clients.doc_cache import InstructionDocCache
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 (
    ENV_IMS_DEBUG,
    ENV_ROSETTA_MODE,
    TAG_BOOTSTRAP,
    TAG_MCP_SERVER_INSTRUCTIONS,
    TOOL_DISCOVER_PROJECTS,
    TOOL_GET_CONTEXT_INSTRUCTIONS,
    TOOL_LIST_INSTRUCTIONS,
    TOOL_QUERY_INSTRUCTIONS,
    TOOL_QUERY_PROJECT_CONTEXT,
    TOOL_STORE_PROJECT_CONTEXT,
    TOOL_SUBMIT_FEEDBACK,
    TRANSPORT_HTTP,
    TRANSPORT_STDIO,
)
from ims_mcp.services.authorizer import Authorizer
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_HARD,
    PROMPT_GET_CONTEXT_INSTRUCTIONS_SOFT,
    PROMPT_LIST_INSTRUCTIONS,
    PROMPT_QUERY_INSTRUCTIONS,
    PROMPT_QUERY_PROJECT_CONTEXT,
    PROMPT_SERVER_INSTRUCTIONS_HARD,
    PROMPT_SERVER_INSTRUCTIONS_SOFT,
    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 list_instructions as list_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,
)
from ims_mcp.tools.resources import read_instruction_resource

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


def _build_session_store():
    if not _CONFIG.redis_url:
        return None  # FastMCP defaults to MemoryStore (local dev / single-instance)
    try:
        from key_value.aio.stores.redis import RedisStore
        return RedisStore(url=_CONFIG.redis_url)
    except ImportError:
        debug_print(_CONFIG, "[ims-mcp] py-key-value-aio[redis] not installed; falling back to in-memory store")
        return None


_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)
_DOC_CACHE = InstructionDocCache(_DOCUMENT_CLIENT)
_FEEDBACK = FeedbackService()
_CONTEXT_INSTRUCTIONS_CACHE: str | None = None
_AUTHORIZER = Authorizer(_CONFIG.read_policy, _CONFIG.write_policy)
_OAUTH_PROVIDER = build_oauth_provider(_CONFIG)
_IS_DEBUG = os.getenv(ENV_IMS_DEBUG, "").strip().lower() in {"1", "true", "yes", "on"}


if _IS_DEBUG:
    print(f"Rosetta v{_MCP_VERSION}", file=sys.stderr, flush=True)


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_mcp_server_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_MCP_SERVER_INSTRUCTIONS]),
        )
        return _BUNDLER.bundle(docs, dataset_name) if docs else ""
    except Exception as exc:
        debug_print(_CONFIG, f"[ims-mcp] server instructions load failed: {exc}")
        return ""


# Intentionally not relied on: many IDE clients do not consistently inject
# server-level instructions into agent context.
#_MCP_SERVER_INSTRUCTIONS = _load_mcp_server_instructions()
_ROSETTA_ICON = load_rosetta_icon()

# Select prompt mode based on ROSETTA_MODE environment variable
# HARD is the default mode
_rosetta_mode = os.getenv(ENV_ROSETTA_MODE, "HARD").upper()

if _rosetta_mode == "SOFT":
    PROMPT_SERVER_INSTRUCTIONS = PROMPT_SERVER_INSTRUCTIONS_SOFT
    PROMPT_GET_CONTEXT_INSTRUCTIONS = PROMPT_GET_CONTEXT_INSTRUCTIONS_SOFT
else:
    # Default to HARD mode
    PROMPT_SERVER_INSTRUCTIONS = PROMPT_SERVER_INSTRUCTIONS_HARD
    PROMPT_GET_CONTEXT_INSTRUCTIONS = PROMPT_GET_CONTEXT_INSTRUCTIONS_HARD

mcp = FastMCP(
    name="Rosetta",
    # Server-level instructions are not relied on for prep/context because
    # many IDE clients do not consistently inject them into agent context.
    instructions=(
        f"Rosetta v{_MCP_VERSION} is a centralized instruction and knowledge "
        "management system intentionally installed by the user for you to follow as a must."
        + PROMPT_SERVER_INSTRUCTIONS
    ),
    # instructions=_MCP_SERVER_INSTRUCTIONS,
    icons=[_ROSETTA_ICON] if _ROSETTA_ICON else [],
    session_state_store=_build_session_store(),
    auth=_OAUTH_PROVIDER,
)


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


def _resolve_user_email() -> str:
    """Return the authenticated user's email.

    HTTP mode: extracted from the OAuth access-token claims.
    STDIO mode (or when no token is available): falls back to
    ``ROSETTA_USER_EMAIL`` env var.
    """
    if _CONFIG.transport == TRANSPORT_HTTP:
        try:
            from fastmcp.server.dependencies import get_access_token
            token = get_access_token()
            if token and getattr(token, "claims", None):
                email = token.claims.get("email", "")
                if email:
                    return email
        except Exception:
            pass
    return _CONFIG.user_email


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,
        user_email=_resolve_user_email(),
        authorizer=_AUTHORIZER,
    )


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


# This is required, as sometimes models hallucinate tags as single string, but we don't want tool contract to be different (as it causes more hallucinations)
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.resource(
    "rosetta://{path*}",
    name="instructions",
    description="Read bundled instruction documents by resource path.",
    mime_type="text/plain",
)
async def _read_resource(path: str, ctx: Context | None = None) -> str:
    if not _RAGFLOW:
        return "Error: ROSETTA_API_KEY is required"
    call_ctx = await _build_call_context("resource_read", {"path": path}, ctx)
    return await read_instruction_resource(
        path=path,
        call_ctx=call_ctx,
        document_client=_DOCUMENT_CLIENT,
        bundler=_BUNDLER,
        doc_cache=_DOC_CACHE,
    )


@mcp.tool(name=TOOL_GET_CONTEXT_INSTRUCTIONS, description=PROMPT_GET_CONTEXT_INSTRUCTIONS)
@track_tool_call
async def get_context_instructions(
    topic: Annotated[str | None, Field(description="Brief intent summary (<=10 words). Does not affect the result, used for logging and tracking.")] = None,
    ctx: Context | None = None,
) -> str:
    if not _RAGFLOW:
        return "Error: ROSETTA_API_KEY is required"
    err = _validate_topic(topic)
    if err:
        return err

    global _CONTEXT_INSTRUCTIONS_CACHE
    if _CONTEXT_INSTRUCTIONS_CACHE is not None:
        return _CONTEXT_INSTRUCTIONS_CACHE

    await _log(ctx, "info", "Loading context instructions")
    call_ctx = await _build_call_context(TOOL_GET_CONTEXT_INSTRUCTIONS, {"topic": topic}, ctx)
    result = await _retry_once(
        lambda: query_instructions_tool(
            call_ctx=call_ctx,
            document_client=_DOCUMENT_CLIENT,
            bundler=_BUNDLER,
            query_builder=_QUERY_BUILDER,
            query=None,
            tags=[TAG_BOOTSTRAP],
            topic=None, # no topic, as it creates too many results and noise
        )
    )
    if result and not result.startswith("Error:"):
        _CONTEXT_INSTRUCTIONS_CACHE = result
    return result


@mcp.tool(name=TOOL_QUERY_INSTRUCTIONS, description=PROMPT_QUERY_INSTRUCTIONS)
@track_tool_call
async def query_instructions(
    query: Annotated[str | None, Field(description="Keyword search text for instruction documents.")] = None,
    tags: Annotated[list[str] | None, Field(description='Known files and families tags ("OR" logic).')] = None,
    topic: Annotated[str | None, Field(description="Optional <=10-word intent used for logging only. Does not affect the result.")] = None,
    ctx: Context | None = None,
) -> str:
    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=None, # no topic, as it creates too many results and noise
        )
    )


@mcp.tool(name=TOOL_LIST_INSTRUCTIONS, description=PROMPT_LIST_INSTRUCTIONS)
@track_tool_call
async def list_instructions(
    path_prefix: Annotated[str, Field(description="Virtual folder path to list immediate children of (e.g. 'skills', 'rules', 'skills/coding-agents-prompt-authoring').")],
    ctx: Context | None = None,
) -> str:
    if not _RAGFLOW:
        return "Error: ROSETTA_API_KEY is required"

    await _log(ctx, "info", f"Listing instructions at {path_prefix}")
    call_ctx = await _build_call_context(
        TOOL_LIST_INSTRUCTIONS,
        {"path_prefix": path_prefix},
        ctx,
    )
    return await _retry_once(
        lambda: list_instructions_tool(
            call_ctx=call_ctx,
            doc_cache=_DOC_CACHE,
            bundler=_BUNDLER,
            path_prefix=path_prefix,
        )
    )


@mcp.tool(name=TOOL_SUBMIT_FEEDBACK, description=PROMPT_SUBMIT_FEEDBACK)
@track_tool_call
async def submit_feedback(
    request_mode: Annotated[str, Field(description='Workflow classification. Examples: "coding.md", "help.md", "research.md", "aqa.md".')],
    feedback: Annotated[dict, Field(description=(
        "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."
    ))],
    ctx: Context | None = None,
) -> str:

    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: Annotated[str, Field(description="Name of the workspace folder name (fallback to the repository name).")],
    query: Annotated[str | None, Field(description="Keyword search text for project context documents.")] = None,
    tags: Annotated[list[str] | None, Field(description="Filter by context tags (string or list).")] = None,
    topic: Annotated[str | None, Field(description="Optional <=10-word intent used for logging only. Does not affect the result.")] = None,
    ctx: Context | None = None,
) -> str:

    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=None, # no topic, as it creates too many results and noise
        )
    )


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

    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: Annotated[str | None, Field(description="Optional search term to filter projects by name.")] = None,
    ctx: Context | None = None,
) -> str:

    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))


class OriginValidationMiddleware:
    """ASGI middleware that validates the Origin header against an allowlist.

    Protects against DNS rebinding attacks. No-op when allowlist is empty.
    """

    def __init__(self, app, allowed_origins: list[str]) -> None:
        self.app = app
        self._allowed = set(allowed_origins)

    async def __call__(self, scope, receive, send) -> None:
        if scope["type"] in {"http", "websocket"} and self._allowed:
            headers = {k.lower(): v for k, v in scope.get("headers", [])}
            origin = headers.get(b"origin", b"").decode()
            if origin and origin not in self._allowed:
                if scope["type"] == "http":
                    await send({"type": "http.response.start", "status": 403, "headers": []})
                    await send({"type": "http.response.body", "body": b"Forbidden", "more_body": False})
                    return
        await self.app(scope, receive, send)


def main() -> None:
    if _CONFIG.transport == TRANSPORT_HTTP:
        from starlette.middleware import Middleware
        middleware: list[Middleware] = []
        if _CONFIG.allowed_origins:
            middleware.append(Middleware(OriginValidationMiddleware, allowed_origins=_CONFIG.allowed_origins))
        mcp.run(
            transport="http",
            host=_CONFIG.http_host,
            port=_CONFIG.http_port,
            stateless_http=False,
            middleware=middleware,
        )
    else:
        mcp.run(transport=TRANSPORT_STDIO)


if __name__ == "__main__":
    main()
