"""Verification script for Rosetta MCP.

This script validates:
1) tool input schemas include parameter descriptions,
2) resource template registration exists,
3) VFS resource reads return bundled content,
4) project lifecycle: discover -> store (force) -> discover -> query -> delete.

Required environment:
- VERSION: release version used for dataset selection and test path candidates
  (for example: r1 or r2)

Optional environment:
- VFS_STRICT: strict mode toggle. Default is strict (1).
  Set VFS_STRICT=0 only for diagnostic runs.

Runtime requirement:
- Network access to configured Rosetta/RAGFlow endpoint and readable instruction dataset
  for the selected VERSION (for example: aia-r1 when VERSION=r1).

Run examples:
  VERSION=r1 python ims-mcp-server/tests/test_mcp_verification.py
  VERSION=r2 python ims-mcp-server/tests/test_mcp_verification.py
"""

import asyncio
import os
import sys
import uuid
from typing import List

# Ensure imports resolve to local ims-mcp-server sources, not an installed package.
SERVER_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if SERVER_ROOT not in sys.path:
    sys.path.insert(0, SERVER_ROOT)

# Load API keys from repo-root tools/.env if not already set
env_file = os.path.join(os.path.dirname(__file__), "..", "..", "tools", ".env")
if os.path.exists(env_file):
    with open(env_file) as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith("#") and "=" in line:
                key, _, value = line.partition("=")
                os.environ.setdefault(key.strip(), value.strip())

from fastmcp import Client
from ims_mcp.server import mcp  # import the FastMCP instance directly


def get_test_resource_paths(version: str) -> List[str]:
    """Return candidate resource_path values for the configured version."""
    if version.startswith("r2"):
        return [
            "rules/requirements-best-practices.md",
            "agents/prompt-engineer.md",
            "workflows/requirements-flow.md",
        ]
    return [
        "coding.md",
        "agents.md",
        "guardrails.md",
    ]


def is_successful_bundle(content: str) -> bool:
    """True when resource read returned bundled XML payload."""
    return "<rosetta:file " in content


def extract_text(result) -> str:
    """Best-effort extraction of text from FastMCP call/read results."""
    if result is None:
        return ""

    # Newer FastMCP call_tool return type: CallToolResult(content=[TextContent...], data=...)
    content = getattr(result, "content", None)
    if isinstance(content, list) and content:
        first = content[0]
        text = getattr(first, "text", None)
        if text is not None:
            return text
        return str(first)

    # Older/list-style return payloads
    if isinstance(result, (list, tuple)) and result:
        first = result[0]
        text = getattr(first, "text", None)
        if text is not None:
            return text
        return str(first)

    text = getattr(result, "text", None)
    if text is not None:
        return text

    data = getattr(result, "data", None)
    if isinstance(data, str):
        return data
    if data is not None:
        return str(data)

    return str(result)


async def main():
    errors = []

    # Connect to the MCP server via in-memory transport
    print("Connecting to MCP server via in-memory transport...")
    async with Client(mcp) as client:
        # 1. Verify tool schemas have parameter descriptions
        print("\n=== Tool Schema Verification ===")
        tools = await client.list_tools()
        print(f"Found {len(tools)} tools")

        for tool in tools:
            print(f"\nTool: {tool.name}")
            schema = tool.inputSchema
            props = schema.get("properties", {})
            for param_name, param_info in props.items():
                desc = param_info.get("description", "")
                status = "OK" if desc else "MISSING"
                if not desc:
                    errors.append(f"Tool '{tool.name}' param '{param_name}' has no description")
                print(f"  {param_name}: [{status}] {desc[:80] if desc else '<no description>'}")

        # 2. Verify resource templates are registered
        print("\n=== Resource Template Verification ===")
        templates = await client.list_resource_templates()
        print(f"Found {len(templates)} resource templates")

        for tmpl in templates:
            print(f"  Template: {tmpl.uriTemplate} (name={tmpl.name})")

        if not templates:
            errors.append("No resource templates registered")
        else:
                # Check we have our rosetta-* template
                rosetta_templates = [t for t in templates if "rosetta" in str(t.uriTemplate)]
                if not rosetta_templates:
                    errors.append("No rosetta-* resource template found")
                else:
                    print(f"\n  Found rosetta template: {rosetta_templates[0].uriTemplate}")
                    if str(rosetta_templates[0].uriTemplate) != "rosetta://{path*}":
                        errors.append(
                            "Unexpected rosetta resource template URI. Expected 'rosetta://{path*}'"
                        )

        # 3. Verify VFS resource read behavior
        print("\n=== VFS Resource Read Verification ===")
        version = (os.environ.get("VERSION") or "").strip()
        if not version:
            errors.append("VERSION environment variable is required (example: VERSION=r1)")
            version = "r1"
        strict_vfs = os.environ.get("VFS_STRICT", "1").lower() not in {"0", "false", "no"}
        candidate_paths = get_test_resource_paths(version)
        bundle_found = False
        no_docs_count = 0
        connection_errors = 0

        for resource_path in candidate_paths:
            test_path = f"rosetta://{resource_path}"
            try:
                result = await client.read_resource(test_path)
                if result and hasattr(result[0], "text"):
                    content = result[0].text
                else:
                    content = str(result)

                print(f"  Read {test_path}: {len(content)} chars")
                print(f"  Preview: {content[:220]}...")

                if is_successful_bundle(content):
                    bundle_found = True
                    print("  Result: BUNDLED CONTENT FOUND")
                    break

                if content.startswith("No documents found for resource path:"):
                    no_docs_count += 1
                    print("  Result: dataset reachable, no matching resource_path")
                elif content.startswith("Error:"):
                    connection_errors += 1
                    print("  Result: resource read error")
                else:
                    connection_errors += 1
                    print("  Result: non-bundled response (unexpected format)")

            except Exception as exc:
                connection_errors += 1
                print(f"  Read {test_path}: {exc}")

        if not bundle_found:
            if strict_vfs:
                errors.append(
                    "VFS strict mode failed: no bundled content found for any candidate path "
                    f"for version {version}"
                )
            else:
                print(
                    f"  NOTE: No bundled content found for version {version}. "
                    f"no_docs={no_docs_count}, errors={connection_errors}. "
                    "Set VFS_STRICT=1 to fail in this case."
                )
        if strict_vfs and connection_errors:
            errors.append(f"VFS resource read had {connection_errors} error response(s)")

        # 4. Project lifecycle: discover -> store (force) -> discover -> query -> store (update) -> query -> delete
        print("\n=== Project Lifecycle Verification ===")
        test_project = f"mcp-verify-{uuid.uuid4().hex[:8]}"
        test_doc = "VERIFY-TEST.md"
        test_content_v1 = "# Verification Test V1\nThis is the initial verification document."
        test_content_v2 = "# Verification Test V2\nThis is the updated verification document."
        test_tags = ["test", "verification"]
        dataset_name = f"project-{test_project}"

        try:
            # Step 1: discover — project must not exist yet
            print(f"\n  Step 1: Discover projects (expect '{test_project}' absent)")
            result = await client.call_tool("discover_projects", {"query": test_project})
            discover_text = extract_text(result)
            print(f"    Result: {discover_text[:200]}")
            if test_project in discover_text and "No projects found" not in discover_text:
                errors.append(f"Project '{test_project}' already exists before test")

            # Step 2: store with force — creates project and document
            print(f"\n  Step 2: Store document with force=true (v1)")
            result = await client.call_tool("store_project_context", {
                "repository_name": test_project,
                "document": test_doc,
                "tags": test_tags,
                "content": test_content_v1,
                "force": True,
            })
            store_text = extract_text(result)
            print(f"    Result: {store_text[:200]}")
            if not store_text.startswith("Stored"):
                errors.append(f"store_project_context (v1) failed: {store_text}")

            # Step 3: discover — project must now appear
            print(f"\n  Step 3: Discover projects (expect '{test_project}' present)")
            result = await client.call_tool("discover_projects", {"query": test_project})
            discover_text = extract_text(result)
            print(f"    Result: {discover_text[:200]}")
            if test_project not in discover_text:
                errors.append(f"Project '{test_project}' not found after store")

            # Step 4: query — v1 document must be retrievable
            print(f"\n  Step 4: Query project context (expect v1)")
            result = await client.call_tool("query_project_context", {
                "repository_name": test_project,
                "tags": test_tags,
            })
            query_text = extract_text(result)
            print(f"    Result: {query_text[:200]}...")
            if not is_successful_bundle(query_text):
                errors.append(f"query_project_context (v1) did not return bundled content: {query_text[:100]}")
            elif "V1" not in query_text:
                errors.append(f"query_project_context (v1) content mismatch: expected V1 marker")

            # Step 5: store again (update) — same document name, new content
            print(f"\n  Step 5: Store document again (update, v2)")
            result = await client.call_tool("store_project_context", {
                "repository_name": test_project,
                "document": test_doc,
                "tags": test_tags,
                "content": test_content_v2,
            })
            store_text = extract_text(result)
            print(f"    Result: {store_text[:200]}")
            if not store_text.startswith("Stored"):
                errors.append(f"store_project_context (v2 update) failed: {store_text}")

            # Step 6: query — must return v2, not v1
            print(f"\n  Step 6: Query project context (expect v2)")
            result = await client.call_tool("query_project_context", {
                "repository_name": test_project,
                "tags": test_tags,
            })
            query_text = extract_text(result)
            print(f"    Result: {query_text[:200]}...")
            if not is_successful_bundle(query_text):
                errors.append(f"query_project_context (v2) did not return bundled content: {query_text[:100]}")
            elif "V2" not in query_text:
                errors.append(f"query_project_context (v2) content mismatch: expected V2 marker, got V1")

        finally:
            # Cleanup: delete the test dataset via SDK
            print(f"\n  Cleanup: Deleting test dataset '{dataset_name}'")
            try:
                from ims_mcp.server import _RAGFLOW
                if _RAGFLOW:
                    datasets = _RAGFLOW.list_datasets(page=1, page_size=1000)
                    for ds in datasets:
                        if ds.name == dataset_name:
                            _RAGFLOW.delete_datasets([ds.id])
                            print(f"    Deleted dataset '{dataset_name}' (id: {ds.id})")
                            break
                    else:
                        print(f"    Dataset '{dataset_name}' not found (already clean)")
                else:
                    print("    WARNING: _RAGFLOW not available, manual cleanup needed")
            except Exception as cleanup_exc:
                print(f"    WARNING: cleanup failed: {cleanup_exc}")

    # Summary
    print("\n=== Summary ===")
    if errors:
        print(f"FAILED: {len(errors)} error(s)")
        for e in errors:
            print(f"  - {e}")
        sys.exit(1)
    else:
        print("ALL CHECKS PASSED")


if __name__ == "__main__":
    asyncio.run(main())
