"""
Agent 4: Helix-QA Framework Writer — MCP Server

Writes files produced by agent_playwright_generator into the correct directory
layout of a local Helix-QA project (Playwright + TypeScript + Cucumber + self-healing).

No LLM calls. No ADO calls. Pure file-system operations.

Four tools:
  inspect_helix_project  — Detect framework state before writing anything.
                           Returns "absent" | "partial" | "present" and a
                           recommended write mode.
  write_helix_files      — Write generated files with full deduplication and
                           interface adaptation.
  read_helix_file        — Read one existing file (pre-write content check).
  list_helix_tree        — List all .ts / .feature files grouped by category.

Skills: see skills/write-helix-files.md
"""
from __future__ import annotations

import asyncio
import json
import re
import sys
import tempfile
import time
from pathlib import Path

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

from stlc_agents.agent_helix_writer.tools.helix_write import (
    inspect_helix_project as _inspect,
    write_files_to_helix  as _write_files,
    read_helix_file       as _read_file,
    list_helix_tree       as _list_tree,
    update_helix_file     as _update_file,
)
from stlc_agents.shared.cost_tracker import track

load_dotenv()

app = Server("qa-helix-writer")


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

def _validate_write_inputs(files: dict) -> dict:
    """Validate file dict before writing to Helix-QA project.

    Checks:
      1. Files dict is non-empty.
      2. Every file has non-empty content.
      3. File keys use consistent naming patterns.
      4. Basic TS syntax (brace/paren balance) for .ts files.

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

    if not files:
        errors.append("No files provided. Pass the 'files' dict from generate_playwright_code.")
        return {"valid": False, "errors": errors, "warnings": warnings}

    for file_key, content in files.items():
        if not content or not content.strip():
            errors.append(f"'{file_key}': content is empty.")
            continue

        # TS syntax balance check
        if file_key.endswith(".ts"):
            s = re.sub(r'`[^`]*`', '``', content)
            s = re.sub(r'"[^"\n]*"', '""', s)
            s = re.sub(r"'[^'\n]*'", "''", s)
            s = re.sub(r'//[^\n]*', '', s)
            s = re.sub(r'/\*.*?\*/', '', s, flags=re.DOTALL)
            brace_delta = s.count('{') - s.count('}')
            if brace_delta != 0:
                errors.append(f"'{file_key}': unbalanced braces (delta={brace_delta:+d})")
            paren_delta = s.count('(') - s.count(')')
            if paren_delta != 0:
                errors.append(
                    f"'{file_key}': unbalanced parentheses (delta={paren_delta:+d}). "
                    f"Cause: the regex-based string preprocessor strips template literals, "
                    f"quoted strings, and comments before counting — any unescaped '(' or ')' "
                    f"inside a raw regex step pattern will be miscounted. "
                    f"Fix: replace raw regex step patterns (e.g. /^I click (.+)$/) with "
                    f"Cucumber expressions (e.g. 'I click {{string}}') so parentheses appear "
                    f"only inside well-formed TypeScript function signatures."
                )
            if re.search(r"""from\s+['"]{2}""", content):
                errors.append(f"'{file_key}': empty import path (from \"\")")

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


def _validate_write_result(result: dict) -> dict:
    """Validate write_files_to_helix result before returning to user.

    Checks:
      1. At least one file was written (or all were intentionally skipped).
      2. No unexpected failures.
      3. Deduplication report is consistent.

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

    if not result.get("success", False) and result.get("error"):
        errors.append(f"Write failed: {result['error']}")
        return {"valid": False, "errors": errors, "warnings": warnings}

    written = result.get("written", [])
    skipped = result.get("skipped", [])
    summary = result.get("summary", {})

    if not written and skipped:
        # All files were skipped — check if intentional
        skip_reasons = [s.get("reason", "") for s in skipped]
        infra_skips = [r for r in skip_reasons if "infrastructure" in r.lower() or "already exists" in r.lower()]
        if len(infra_skips) == len(skipped):
            warnings.append("All files were skipped (already exist or infrastructure-only). No new content written.")
        else:
            error_skips = [s for s in skipped if "error" in s.get("reason", "").lower() or "empty" in s.get("reason", "").lower()]
            if error_skips:
                errors.append(
                    f"{len(error_skips)} file(s) skipped due to errors: "
                    + ", ".join(s.get("file_key", "?") for s in error_skips)
                )

    # Check for write count mismatches
    if summary.get("requested", 0) > 0:
        write_rate = summary.get("written", 0) / summary["requested"]
        if write_rate == 0 and not skipped:
            errors.append("No files were written and none were skipped — unexpected state.")
        elif write_rate < 0.5 and summary["requested"] > 2:
            warnings.append(
                f"Only {summary.get('written', 0)}/{summary['requested']} files were written. "
                f"Check skipped files for issues."
            )

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


def _validate_inspect_result(result: dict) -> dict:
    """Validate inspect_helix_project result before returning to user.

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

    state = result.get("framework_state", "")
    if state not in ("absent", "partial", "present"):
        errors.append(f"Unknown framework_state: '{state}'")

    if state == "absent":
        warnings.append(
            "Helix-QA framework is absent. Run scaffold_locator_repository "
            "before writing test files."
        )
    elif state == "partial":
        missing = result.get("missing_infra", [])
        warnings.append(
            f"Helix-QA framework is partially set up. Missing: {missing}. "
            f"Run scaffold_locator_repository to complete the setup."
        )

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


def _validate_read_file_response(result: dict) -> dict:
    """Validate read_helix_file result before returning to user.

    Checks:
      1. No error in response.
      2. File content was actually returned.
      3. Warns if content is suspiciously short.

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

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

    content = result.get("content", "")
    if not content and not result.get("found", True):
        warnings.append("File not found on disk. It may not exist yet.")
    elif not content:
        warnings.append("File exists but has no content (empty file).")
    elif len(content.strip().splitlines()) < 3:
        warnings.append(
            "File content is very short (fewer than 3 lines). "
            "It may be a stub or incomplete."
        )

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


def _validate_list_tree_response(result: dict) -> dict:
    """Validate list_helix_tree result before returning to user.

    Checks:
      1. No error in response.
      2. At least one category has files (warns if completely empty).

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

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

    categories = ["features", "steps", "pages", "locators", "utils_locators"]
    total_files = 0
    for cat in categories:
        total_files += len(result.get(cat, []))

    if total_files == 0:
        warnings.append(
            "No .ts or .feature files found in the Helix-QA project. "
            "The project may be empty or the helix_root path may be incorrect."
        )

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


def _validate_update_file_result(result: dict) -> dict:
    """Validate update_helix_file result before returning to user.

    Checks:
      1. No error in response.
      2. File was actually written (non-zero bytes).
      3. Warns when all merge candidates were duplicates (nothing new added).

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

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

    if result.get("bytes", 0) == 0:
        warnings.append("File was written with zero bytes — the content may have been empty after adaptation.")

    dedup = result.get("deduplication") or {}
    added = dedup.get("added_keys") or dedup.get("added_patterns") or dedup.get("added_methods") or []
    skipped = dedup.get("skipped_keys") or dedup.get("skipped_patterns") or dedup.get("skipped_methods") or []
    if skipped and not added:
        warnings.append(
            f"All {len(skipped)} item(s) in the provided content already exist in the file — "
            "nothing new was merged. If you intended to replace existing entries, "
            "call update_helix_file with force_overwrite=true."
        )

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


@app.list_tools()
async def list_tools() -> list[types.Tool]:
    return [
        types.Tool(
            name="inspect_helix_project",
            description=(
                "Detect the state of a local Helix-QA project before writing any files. "
                "Returns framework_state ('absent' | 'partial' | 'present'), "
                "which infrastructure files exist or are missing, and a recommended "
                "write mode ('scaffold_and_tests' or 'tests_only'). "
                "Always call this as Step 1 — the result determines what write_helix_files should do:\n"
                "  'absent' / 'partial' → run scaffold_locator_repository first, "
                "then write_helix_files with mode='scaffold_and_tests'.\n"
                "  'present' → call write_helix_files with mode='tests_only' "
                "(infrastructure files are never touched)."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "helix_root": {
                        "type": "string",
                        "description": "Absolute path to the Helix-QA project root (contains package.json and src/).",
                    },
                },
                "required": ["helix_root"],
            },
        ),
        types.Tool(
            name="write_helix_files",
            description=(
                "Write generated TypeScript/Gherkin files into the Helix-QA directory layout "
                "with full deduplication and interface adaptation.\n\n"
                "Preferred: pass cache_key from generate_playwright_code — "
                "the server loads the files from disk automatically. "
                "Alternative: pass the 'files' dict directly.\n\n"
                "mode='tests_only' (default, safe to run repeatedly):\n"
                "  Writes locators.ts, *.page.ts, *.steps.ts, *.feature only.\n"
                "  Infrastructure files (LocatorHealer.ts etc.) are always skipped.\n"
                "  Existing files are MERGED — new locator keys / step patterns / "
                "async methods are appended; duplicates are skipped.\n\n"
                "mode='scaffold_and_tests':\n"
                "  Also writes infrastructure files, but only if they do not yet exist.\n"
                "  Use when inspect_helix_project returns 'absent' or 'partial'.\n\n"
                "force_scaffold=true:\n"
                "  Overwrite existing infrastructure files. Use only when deliberately "
                "regenerating the healing infrastructure.\n\n"
                "Interface adaptation is applied automatically:\n"
                "  repo.updateHealed → repo.setHealed\n"
                "  repo.incrementSuccess / incrementFailure / queueSuggestion → removed\n"
                "  repo.getBBox → null\n"
                "  fixture().logger / .page / .locatorRepository → constructor arguments\n"
                "  EnvironmentManager → Helix environment singleton\n"
                "  Winston Logger → HealerLogger"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "helix_root": {
                        "type": "string",
                        "description": "Absolute path to the Helix-QA project root.",
                    },
                    "cache_key": {
                        "type": "string",
                        "description": (
                            "Cache key returned by generate_playwright_code or "
                            "scaffold_locator_repository. Preferred over passing 'files' directly — "
                            "the server loads the files from disk automatically."
                        ),
                    },
                    "files": {
                        "type": "object",
                        "description": (
                            "Dict of { file_key: file_content } as returned by "
                            "generate_playwright_code or scaffold_locator_repository. "
                            "Use cache_key instead whenever possible."
                        ),
                        "additionalProperties": {"type": "string"},
                    },
                    "mode": {
                        "type": "string",
                        "enum": ["tests_only", "scaffold_and_tests"],
                        "description": (
                            "Write mode. Use 'tests_only' when the framework already exists. "
                            "Use 'scaffold_and_tests' when inspect_helix_project returns "
                            "'absent' or 'partial'. Default: 'tests_only'."
                        ),
                    },
                    "force_scaffold": {
                        "type": "boolean",
                        "description": (
                            "When true and mode='scaffold_and_tests', overwrite existing "
                            "infrastructure files. Default: false."
                        ),
                    },
                },
                "required": ["helix_root"],
            },
        ),
        types.Tool(
            name="read_helix_file",
            description=(
                "Read the content of one file from the Helix-QA project. "
                "Use before writing to check whether a locators, page object, or steps file "
                "already exists — and what locator keys / step patterns / methods it contains — "
                "so you can confirm the merge will not lose hand-edited content."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "helix_root": {
                        "type": "string",
                        "description": "Absolute path to the Helix-QA project root.",
                    },
                    "relative_path": {
                        "type": "string",
                        "description": (
                            "Path relative to helix_root, e.g. "
                            "'src/locators/login-page.locators.ts'."
                        ),
                    },
                },
                "required": ["helix_root", "relative_path"],
            },
        ),
        types.Tool(
            name="list_helix_tree",
            description=(
                "List all .ts and .feature files in the Helix-QA project, grouped by category "
                "(features, steps, pages, locators, utils_locators). "
                "Call this after inspect_helix_project to see exactly which per-feature files "
                "already exist before deciding what write_helix_files needs to produce."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "helix_root": {
                        "type": "string",
                        "description": "Absolute path to the Helix-QA project root.",
                    },
                },
                "required": ["helix_root"],
            },
        ),
        types.Tool(
            name="update_helix_file",
            description=(
                "Write or merge a single file in the Helix-QA project without touching any other files. "
                "Use this for targeted single-file edits — e.g. adding one missing locator key, "
                "fixing one step pattern, or correcting one page method — instead of re-submitting "
                "the entire files dict to write_helix_files.\n\n"
                "Applies the same Helix interface adapter and merge logic as write_helix_files:\n"
                "  locators.ts  → new const-object entries appended; duplicate keys skipped.\n"
                "  *.steps.ts   → new step blocks appended; duplicate regex patterns skipped.\n"
                "  *.page.ts    → new async methods appended; duplicate method names skipped.\n"
                "  *.feature    → always overwritten (Gherkin is the source of truth).\n"
                "  other .ts    → merged or overwritten depending on force_overwrite.\n\n"
                "Set force_overwrite=true to replace the entire file without merging "
                "(use only when you intentionally want to discard existing hand-edited content)."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "helix_root": {
                        "type": "string",
                        "description": "Absolute path to the Helix-QA project root.",
                    },
                    "relative_path": {
                        "type": "string",
                        "description": (
                            "Path of the file to write, relative to helix_root. "
                            "E.g. 'src/locators/login.locators.ts' or "
                            "'src/test/steps/login.steps.ts'."
                        ),
                    },
                    "content": {
                        "type": "string",
                        "description": "Full or partial TypeScript / Gherkin content to write or merge into the file.",
                    },
                    "force_overwrite": {
                        "type": "boolean",
                        "description": (
                            "When true, replace the entire file without merging. "
                            "Default: false (merge / append only)."
                        ),
                    },
                },
                "required": ["helix_root", "relative_path", "content"],
            },
        ),
    ]


@app.call_tool()
async def call_tool(name: str, arguments: dict) -> list[types.TextContent]:
    t0 = time.monotonic()
    try:
        if name == "inspect_helix_project":
            result = await asyncio.to_thread(_inspect, arguments["helix_root"])
            # ── Pre-output validation ─────────────────────────────────────
            result["_validation"] = _validate_inspect_result(result)

        elif name == "write_helix_files":
            # ── Resolve files: cache_key takes priority over inline dict ──
            files = arguments.get("files") or {}
            # Normalise: LLMs sometimes send [{file_name, content}] instead of {path: content}
            if isinstance(files, list):
                files = {
                    item.get("file_name") or item.get("path") or item.get("name", ""): item.get("content", "")
                    for item in files
                    if isinstance(item, dict)
                }
            cache_key = arguments.get("cache_key", "").strip()
            if cache_key and not files:
                _cache_dir = Path(tempfile.gettempdir()) / "stlc_file_cache"
                cache_file = _cache_dir / f"{cache_key}.json"
                if cache_file.exists():
                    files = json.loads(cache_file.read_text())
                else:
                    result = {
                        "success": False,
                        "error": f"cache_key '{cache_key}' not found — file does not exist at {cache_file}. "
                                 "Either pass the files dict directly or call get_generated_files first.",
                        "_validation": {"valid": False, "errors": [f"cache_key '{cache_key}' not found"], "warnings": []},
                    }
                    return track(result, tool_name=name, server="qa-helix-writer", t0=t0)
            # ── Pre-write input validation ────────────────────────────────
            input_validation = _validate_write_inputs(files)
            if not input_validation["valid"]:
                result = {
                    "success": False,
                    "error": "input_validation_failed",
                    "_validation": input_validation,
                    "message": (
                        "File inputs failed pre-write validation. Fix the errors "
                        "below and retry. No files were written to disk."
                    ),
                }
                return track(result, tool_name=name, server="qa-helix-writer", t0=t0)

            result = await asyncio.to_thread(
                _write_files,
                arguments["helix_root"],
                files,
                arguments.get("mode", "tests_only"),
                arguments.get("force_scaffold", False),
            )
            # ── Post-write validation ─────────────────────────────────────
            result["_validation"] = _validate_write_result(result)

        elif name == "read_helix_file":
            result = await asyncio.to_thread(
                _read_file,
                arguments["helix_root"],
                arguments["relative_path"],
            )
            # ── Pre-output validation ─────────────────────────────────────
            result["_validation"] = _validate_read_file_response(result)

        elif name == "list_helix_tree":
            result = await asyncio.to_thread(_list_tree, arguments["helix_root"])
            # ── Pre-output validation ─────────────────────────────────────
            result["_validation"] = _validate_list_tree_response(result)

        elif name == "update_helix_file":
            result = await asyncio.to_thread(
                _update_file,
                arguments["helix_root"],
                arguments["relative_path"],
                arguments["content"],
                arguments.get("force_overwrite", False),
            )
            # ── Post-write validation ─────────────────────────────────────
            result["_validation"] = _validate_update_file_result(result)

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

        return track(result, tool_name=name, server="qa-helix-writer", t0=t0)

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


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


def main():
    asyncio.run(_run())


if __name__ == "__main__":
    main()
