"""
Agent 2: QA Gherkin Generator — MCP Server

Tools:
  fetch_feature_hierarchy       — Fetch Feature + all child PBIs + linked TCs with steps
  fetch_work_item_for_gherkin   — Fetch a single PBI or Bug ready for scoped Gherkin (NEW)
  attach_gherkin_to_feature     — Upload and attach a .feature file to a Feature work item
  attach_gherkin_to_work_item   — Upload and attach a scoped .feature file to a PBI/Bug (NEW)
  validate_gherkin_content      — Structurally validate a .feature file before attaching (NEW)

Authentication:
  Silent from ~/.msal-cache/msal-cache.json (shared with microsoft/azure-devops-mcp).
  Browser opens once only if no cached token is found.

Skills: see skills/generate-gherkin.md
"""
import asyncio
import json
import sys
import time

from dotenv import load_dotenv
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp import types

from stlc_agents.shared.auth import get_auth_headers, get_signed_in_user
from stlc_agents.agent_gherkin_generator.tools.ado_gherkin import (
    fetch_feature_hierarchy as _fetch_feature,
    fetch_work_item_for_gherkin as _fetch_wi_for_gherkin,
    attach_feature_file as _attach_feature,
    attach_work_item_file as _attach_wi_file,
    validate_gherkin_content as _validate_gherkin,
)
from stlc_agents.shared.cost_tracker import track

load_dotenv()

app = Server("qa-gherkin-generator")


# ---------------------------------------------------------------------------
# Pre-output validation helpers
# ---------------------------------------------------------------------------

def _validate_hierarchy_response(result: dict, scope: str) -> dict:
    """Validate fetch_feature_hierarchy / fetch_work_item_for_gherkin
    response quality before returning to the user.

    Checks:
      1. Response contains no error key.
      2. Feature / work item has a title.
      3. Acceptance criteria exists (warns if missing).
      4. Child PBIs/Bugs exist for Feature/Epic scope (warns if none).
      5. Test case coverage context is present.

    Returns { valid: bool, errors: list[str], warnings: list[str], gaps: list[str] }.
    """
    errors: list[str] = []
    warnings: list[str] = []
    gaps: list[str] = []

    if "error" in result:
        return {"valid": True, "errors": [], "warnings": [], "gaps": []}

    if scope == "epic":
        epic = result.get("epic", {})
        if not epic.get("title"):
            errors.append("Epic has no title — ADO data may be incomplete.")
        features = result.get("features", [])
        if not features:
            gaps.append(
                "Epic has no child Features. Cannot generate Gherkin without "
                "Feature-level scope. Ask the user to verify the Epic ID."
            )
        for i, feat_data in enumerate(features):
            feat = feat_data.get("feature", {})
            if not feat.get("acceptance_criteria"):
                feat_title = feat.get("title") or "#{0}".format(feat.get("id", i + 1))
                warnings.append(
                    "Feature '{}' has no acceptance criteria — generated Gherkin may lack detail.".format(feat_title)
                )
            if not feat_data.get("child_work_items"):
                warnings.append(
                    f"Feature '{feat.get('title', '')}' has no child PBIs/Bugs."
                )
    elif scope == "feature":
        feat = result.get("feature", {})
        if not feat.get("title"):
            errors.append("Feature has no title — ADO data may be incomplete.")
        if not feat.get("acceptance_criteria"):
            gaps.append(
                "Feature has no acceptance criteria. Gherkin generation requires "
                "acceptance criteria to produce meaningful scenarios. "
                "Ask the user to provide AC or check the work item in ADO."
            )
        children = result.get("child_work_items", [])
        if not children:
            warnings.append(
                "Feature has no child PBIs/Bugs. The generated Gherkin will be "
                "based solely on the Feature-level acceptance criteria."
            )
        for child in children:
            if not child.get("acceptance_criteria"):
                warnings.append(
                    f"PBI/Bug '{child.get('title', '')}' (#{child.get('id', '?')}) "
                    f"has no acceptance criteria."
                )
    elif scope == "work_item":
        wi = result.get("work_item", {})
        if not wi.get("title"):
            errors.append("Work item has no title — ADO data may be incomplete.")
        if not wi.get("acceptance_criteria"):
            gaps.append(
                "Work item has no acceptance criteria. Gherkin generation requires "
                "acceptance criteria to produce meaningful scenarios. "
                "Ask the user to provide AC or check the work item in ADO."
            )

    tc_count = result.get("existing_test_cases_count", 0)
    if tc_count > 0:
        warnings.append(
            f"{tc_count} existing test case(s) found. Review them for "
            f"deduplication before generating new Gherkin scenarios."
        )

    return {
        "valid": len(errors) == 0 and len(gaps) == 0,
        "errors": errors,
        "warnings": warnings,
        "gaps": gaps,
    }


def _validate_attach_response(result: dict, scope: str) -> dict:
    """Validate the response from attach_gherkin_to_feature / attach_gherkin_to_work_item
    before returning to the user.

    Checks:
      1. Response indicates success (no error key).
      2. Attachment URL is present.
      3. File name is present and non-empty.

    Returns { valid: bool, errors: list[str], warnings: list[str] }.
    """
    errors: list[str] = []
    warnings: list[str] = []

    if "error" in result:
        errors.append("Attach failed: {}".format(result["error"]))
        return {"valid": False, "errors": errors, "warnings": warnings}

    if not result.get("success", False):
        errors.append("Attach returned success=false without an error message.")
        return {"valid": False, "errors": errors, "warnings": warnings}

    if not result.get("attachment_url"):
        warnings.append(
            "Attachment URL is missing from the response. The file may have "
            "been attached, but the URL was not returned by ADO."
        )

    if not result.get("file_name"):
        warnings.append("File name is missing from the attach response.")

    return {"valid": len(errors) == 0, "errors": errors, "warnings": warnings}


@app.list_tools()
async def list_tools() -> list[types.Tool]:
    return [
        # ── Feature-scoped fetch ─────────────────────────────────────────────
        types.Tool(
            name="fetch_feature_hierarchy",
            description=(
                "Fetch a Feature work item from Azure DevOps along with all child PBIs/Bugs "
                "and their linked test cases. Use this as Step 1 when the user provides a "
                "Feature ID. Generates 5–10 scenarios covering the full Feature. "
                "For a single PBI or Bug ID, use fetch_work_item_for_gherkin instead. "
                "For an Epic ID: do NOT call any fetch tool — instead tell the user: "
                "'Epic-scoped Gherkin generation is not supported. Please provide a Feature ID "
                "or PBI/Bug ID. To find child Features, look up the Epic in Azure DevOps.'"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "feature_id": {"type": "integer", "description": "ADO Feature work item ID"},
                    "organization_url": {"type": "string", "description": "e.g. https://dev.azure.com/myorg"},
                    "project_name": {"type": "string", "description": "ADO project name"},
                },
                "required": ["feature_id", "organization_url", "project_name"],
            },
        ),

        # ── PBI/Bug-scoped fetch ─────────────────────────────────────────────
        types.Tool(
            name="fetch_work_item_for_gherkin",
            description=(
                "Fetch a single PBI or Bug from Azure DevOps with all fields needed to generate "
                "a scoped Gherkin feature file for that work item alone. "
                "Use this when the user provides a PBI or Bug ID and wants a focused .feature "
                "file covering only that work item's acceptance criteria (3–9 scenarios). "
                "Returns the work item's AC, linked test case titles, and parent Feature metadata. "
                "For Feature IDs, use fetch_feature_hierarchy instead. "
                "For Epic IDs, do not call this tool — tell the user to provide a Feature or PBI ID."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "work_item_id": {
                        "type": "integer",
                        "description": "ADO PBI or Bug work item ID",
                    },
                    "organization_url": {
                        "type": "string",
                        "description": "e.g. https://dev.azure.com/myorg",
                    },
                    "project_name": {
                        "type": "string",
                        "description": "ADO project name",
                    },
                },
                "required": ["work_item_id", "organization_url", "project_name"],
            },
        ),

        # ── Existing: Feature-scoped attach (now gated by validation) ────────
        types.Tool(
            name="attach_gherkin_to_feature",
            description=(
                "Validate, upload, and attach a completed .feature file to a Feature work item "
                "in ADO. Structural validation runs automatically before upload — if the file "
                "is missing @smoke or @regression tags, has fewer than 5 or more than 10 "
                "scenarios, contains duplicate scenario titles, or has any scenario without a "
                "When step, the upload is blocked and errors are returned for correction. "
                "The file is named {feature_id}_{feature_title}_regression.feature automatically. "
                "Call this as the final step after generating and reviewing the Gherkin content."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "feature_id": {"type": "integer", "description": "ADO Feature work item ID"},
                    "feature_title": {"type": "string", "description": "Feature title (used in filename)"},
                    "organization_url": {"type": "string"},
                    "project_name": {"type": "string"},
                    "gherkin_content": {
                        "type": "string",
                        "description": "Complete .feature file content as a string",
                    },
                },
                "required": [
                    "feature_id", "feature_title", "organization_url",
                    "project_name", "gherkin_content",
                ],
            },
        ),

        # ── NEW: PBI/Bug-scoped attach ───────────────────────────────────────
        types.Tool(
            name="attach_gherkin_to_work_item",
            description=(
                "Validate, upload, and attach a scoped .feature file directly to a PBI or Bug "
                "work item in ADO. Use this after fetch_work_item_for_gherkin when the user "
                "wants the feature file attached to the PBI/Bug rather than its parent Feature. "
                "Structural validation runs automatically before upload — if the file is missing "
                "@smoke or @regression tags, has fewer than 3 or more than 9 scenarios "
                "(work-item scope limit), contains duplicate titles, or has any scenario without "
                "a When step, the upload is blocked and errors are returned for correction. "
                "The file is named {work_item_id}_{title_kebab}.feature automatically."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "work_item_id": {
                        "type": "integer",
                        "description": "ADO PBI or Bug work item ID",
                    },
                    "work_item_title": {
                        "type": "string",
                        "description": "Work item title (used in filename)",
                    },
                    "organization_url": {"type": "string"},
                    "project_name": {"type": "string"},
                    "gherkin_content": {
                        "type": "string",
                        "description": "Complete .feature file content as a string",
                    },
                },
                "required": [
                    "work_item_id", "work_item_title", "organization_url",
                    "project_name", "gherkin_content",
                ],
            },
        ),

        # ── NEW: Standalone validation ───────────────────────────────────────
        types.Tool(
            name="validate_gherkin_content",
            description=(
                "Structurally validate a .feature file without uploading it to ADO. "
                "Use this to check a file before calling attach_gherkin_to_feature or "
                "attach_gherkin_to_work_item, or to verify hand-edited Gherkin. "
                "Checks: @smoke tag present, @regression tag present, scenario count within "
                "allowed range, every scenario has a When step, no duplicate scenario titles. "
                "Pass scope='feature' for Feature-scoped files (5–10 scenarios) or "
                "scope='work_item' for PBI/Bug-scoped files (3–9 scenarios)."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "gherkin_content": {
                        "type": "string",
                        "description": "Complete .feature file content as a string",
                    },
                    "scope": {
                        "type": "string",
                        "enum": ["feature", "work_item"],
                        "description": (
                            "Scope determines the allowed scenario count range. "
                            "'feature' = 5–10 scenarios, 'work_item' = 3–9 scenarios. "
                            "Default: 'feature'."
                        ),
                    },
                },
                "required": ["gherkin_content"],
            },
        ),

        # ── NEW: Headless composite — fetch + validate + attach with retry ───
        types.Tool(
            name="generate_and_attach_gherkin",
            description=(
                "Headless composite tool: fetch work item data, validate Gherkin content, "
                "and attach to ADO — retrying once internally if validation fails. "
                "Use this from the pipeline or CI runner instead of calling "
                "validate_gherkin_content + attach_gherkin_to_feature/work_item separately. "
                "Accepts pre-generated gherkin_content from the LLM caller. "
                "If validation fails on the first pass the errors are returned as "
                "validation_failed (no retry is possible without a new LLM call — "
                "the pipeline must re-generate and call this tool again). "
                "Returns {status, gherkin_content, attached} on success, "
                "or {status: validation_failed, errors, gherkin_content} on failure."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "work_item_id": {
                        "type": "integer",
                        "description": "ADO Feature, PBI, or Bug ID",
                    },
                    "work_item_title": {
                        "type": "string",
                        "description": "Title of the work item (used in filename)",
                    },
                    "organization_url": {"type": "string"},
                    "project_name": {"type": "string"},
                    "gherkin_content": {
                        "type": "string",
                        "description": "Pre-generated .feature file content from LLM caller",
                    },
                    "scope": {
                        "type": "string",
                        "enum": ["feature", "work_item"],
                        "description": (
                            "'feature' attaches to a Feature work item (5–10 scenarios). "
                            "'work_item' attaches to a PBI or Bug (3–9 scenarios). "
                            "Default: 'work_item'."
                        ),
                    },
                },
                "required": [
                    "work_item_id", "work_item_title", "organization_url",
                    "project_name", "gherkin_content",
                ],
            },
        ),
    ]


@app.call_tool()
async def call_tool(name: str, arguments: dict) -> list[types.TextContent]:
    t0 = time.monotonic()
    try:
        if name == "fetch_feature_hierarchy":
            result = await asyncio.to_thread(
                _fetch_feature,
                arguments["organization_url"],
                arguments["project_name"],
                arguments["feature_id"],
            )
            # ── Pre-output validation ─────────────────────────────────────
            result["_validation"] = _validate_hierarchy_response(result, "feature")

        elif name == "fetch_work_item_for_gherkin":
            result = await asyncio.to_thread(
                _fetch_wi_for_gherkin,
                arguments["organization_url"],
                arguments["project_name"],
                arguments["work_item_id"],
            )
            # ── Pre-output validation ─────────────────────────────────────
            result["_validation"] = _validate_hierarchy_response(result, "work_item")

        elif name == "attach_gherkin_to_feature":
            # ── Pre-attach validation gate ────────────────────────────────
            pre_check = await asyncio.to_thread(
                _validate_gherkin,
                arguments["gherkin_content"],
                "feature",
            )
            if not pre_check["valid"]:
                result = {
                    "success": False,
                    "reason": "pre_attach_validation_failed",
                    "_validation": pre_check,
                    "message": (
                        "Gherkin content failed pre-attach validation. "
                        "Fix the errors below before attaching to ADO."
                    ),
                }
                return track(result, tool_name=name, server="qa-gherkin-generator", t0=t0)

            result = await asyncio.to_thread(
                _attach_feature,
                arguments["organization_url"],
                arguments["project_name"],
                arguments["feature_id"],
                arguments["feature_title"],
                arguments["gherkin_content"],
            )
            # ── Post-attach validation ────────────────────────────────────
            result["_validation"] = _validate_attach_response(result, "feature")

        elif name == "attach_gherkin_to_work_item":
            # ── Pre-attach validation gate ────────────────────────────────
            pre_check = await asyncio.to_thread(
                _validate_gherkin,
                arguments["gherkin_content"],
                "work_item",
            )
            if not pre_check["valid"]:
                result = {
                    "success": False,
                    "reason": "pre_attach_validation_failed",
                    "_validation": pre_check,
                    "message": (
                        "Gherkin content failed pre-attach validation. "
                        "Fix the errors below before attaching to ADO."
                    ),
                }
                return track(result, tool_name=name, server="qa-gherkin-generator", t0=t0)

            result = await asyncio.to_thread(
                _attach_wi_file,
                arguments["organization_url"],
                arguments["project_name"],
                arguments["work_item_id"],
                arguments["work_item_title"],
                arguments["gherkin_content"],
            )
            # ── Post-attach validation ────────────────────────────────────
            result["_validation"] = _validate_attach_response(result, "work_item")

        elif name == "validate_gherkin_content":
            result = await asyncio.to_thread(
                _validate_gherkin,
                arguments["gherkin_content"],
                arguments.get("scope", "feature"),
            )

        elif name == "generate_and_attach_gherkin":
            org = arguments["organization_url"]
            project = arguments["project_name"]
            wi_id = arguments["work_item_id"]
            wi_title = arguments["work_item_title"]
            gherkin = arguments["gherkin_content"]
            scope = arguments.get("scope", "work_item")

            # Step 1: validate
            validation = await asyncio.to_thread(_validate_gherkin, gherkin, scope)
            if not validation["valid"]:
                result = {
                    "status": "validation_failed",
                    "errors": validation.get("errors", []),
                    "gherkin_content": gherkin,
                    "message": (
                        "Gherkin content failed structural validation. "
                        "The pipeline must re-generate and call this tool again "
                        "with corrected content. No file was attached to ADO."
                    ),
                }
            else:
                # Step 2: attach
                if scope == "feature":
                    attach = await asyncio.to_thread(
                        _attach_feature, org, project, wi_id, wi_title, gherkin
                    )
                else:
                    attach = await asyncio.to_thread(
                        _attach_wi_file, org, project, wi_id, wi_title, gherkin
                    )
                attach["_validation"] = _validate_attach_response(attach, scope)
                result = {
                    "status": "ok",
                    "gherkin_content": gherkin,
                    "attached": attach,
                }

        else:
            result = {"error": f"Unknown tool: {name}"}

        return track(result, tool_name=name, server="qa-gherkin-generator", t0=t0)

    except Exception as exc:
        err_result = {"error": str(exc), "tool": name}
        return track(err_result, tool_name=name, server="qa-gherkin-generator", t0=t0)


async def _run():
    async with stdio_server() as (r, w):
        await app.run(r, w, app.create_initialization_options())


def main():
    try:
        get_auth_headers()
    except Exception as e:
        print(f"[qa-gherkin-generator] Auth error: {e}", file=sys.stderr)
        sys.exit(1)

    user = get_signed_in_user()
    if user:
        print(f"[qa-gherkin-generator] Authenticated as: {user}", file=sys.stderr)

    asyncio.run(_run())


if __name__ == "__main__":
    main()