from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from .engine import assess, build_summary
from .init_tool import init_security_context
from .validation import SecurityContextValidationError, load_and_validate_context, validate_context_file

try:
    from mcp.server.fastmcp import FastMCP
except Exception as exc:  # pragma: no cover - exercised only when mcp is missing.
    FastMCP = None
    MCP_IMPORT_ERROR = exc
else:
    MCP_IMPORT_ERROR = None


@dataclass
class AppState:
    config_path: Path
    schema_path: Path | None
    context: dict[str, Any]


_STATE: AppState | None = None


if FastMCP is not None:
    mcp = FastMCP("aegis")
else:  # pragma: no cover
    mcp = None


def _state() -> AppState:
    if _STATE is None:
        raise RuntimeError("Aegis server state is not initialized.")
    return _STATE


def _tool_aegis_assess(
    services: list[str] | None = None,
    data_types: list[str] | None = None,
    boundary_crossing: str | None = None,
    auth_scope: str | None = None,
    input_sources: list[str] | None = None,
    action: str | None = None,
    description: str | None = None,
) -> dict[str, Any]:
    state = _state()
    return assess(
        context=state.context,
        services=services,
        data_types=data_types,
        boundary_crossing=boundary_crossing,
        auth_scope=auth_scope,
        input_sources=input_sources,
        action=action,
        description=description,
    )


def _tool_aegis_validate(path: str | None = None) -> dict[str, Any]:
    state = _state()
    target = Path(path) if path else state.config_path
    report = validate_context_file(
        context_path=target,
        schema_path=state.schema_path,
    )
    return {
        "valid": report.valid,
        "errors": report.errors,
        "warnings": report.warnings,
        "context_path": str(target),
    }


def _tool_aegis_init(
    threat_model_path: str,
    output_path: str = "security-context.yaml",
    project_name: str | None = None,
    use_llm: bool = True,
    llm_model: str = "claude-3-5-sonnet-latest",
) -> dict[str, Any]:
    state = _state()
    return init_security_context(
        threat_model_path=threat_model_path,
        output_path=output_path,
        project_name=project_name,
        use_llm=use_llm,
        llm_model=llm_model,
        schema_path=state.schema_path,
    )


def _resource_aegis_summary() -> dict[str, Any]:
    state = _state()
    return build_summary(state.context)


if mcp is not None:

    @mcp.tool()
    def aegis_assess(
        services: list[str] | None = None,
        data_types: list[str] | None = None,
        boundary_crossing: str | None = None,
        auth_scope: str | None = None,
        input_sources: list[str] | None = None,
        action: str | None = None,
        description: str | None = None,
    ) -> dict[str, Any]:
        """IMPORTANT: Call BEFORE generating security-relevant code.

        Trigger this tool when the task involves any of:
        - New API endpoints or route handlers
        - Database access or query changes
        - Authentication or authorization logic
        - External API or third-party service integrations
        - File upload or file-processing logic
        - Handling data classified in security-context.yaml

        Pass canonical IDs from security-context.yaml for boundary_crossing,
        data_types, services, and action.

        Returned requirements are mandatory constraints on generated code.
        Prefer approved_patterns, avoid anti_patterns, and if no_match is true
        or warnings include unknown IDs, suggest updating the security context.

        Zero LLM cost: this is a deterministic structured lookup.
        """

        return _tool_aegis_assess(
            services=services,
            data_types=data_types,
            boundary_crossing=boundary_crossing,
            auth_scope=auth_scope,
            input_sources=input_sources,
            action=action,
            description=description,
        )

    @mcp.tool()
    def aegis_validate(path: str | None = None) -> dict[str, Any]:
        """Validate a security-context YAML file against the Aegis schema."""

        return _tool_aegis_validate(path=path)

    @mcp.tool()
    def aegis_init(
        threat_model_path: str,
        output_path: str = "security-context.yaml",
        project_name: str | None = None,
        use_llm: bool = True,
        llm_model: str = "claude-3-5-sonnet-latest",
    ) -> dict[str, Any]:
        """Generate a security-context.yaml from threat-model documentation."""

        return _tool_aegis_init(
            threat_model_path=threat_model_path,
            output_path=output_path,
            project_name=project_name,
            use_llm=use_llm,
            llm_model=llm_model,
        )

    @mcp.resource("aegis://summary")
    def aegis_summary() -> dict[str, Any]:
        """Read once at session start to bootstrap security context.

        Provides trust boundaries, data classifications, and canonical IDs that
        the agent should use when calling aegis_assess.
        """

        return _resource_aegis_summary()


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run the Aegis MCP server")
    parser.add_argument(
        "--config",
        default="security-context.yaml",
        help="Path to security-context.yaml",
    )
    parser.add_argument(
        "--schema",
        default=None,
        help="Optional path to security-context schema file",
    )
    parser.add_argument(
        "--transport",
        default="stdio",
        choices=["stdio", "sse"],
        help="MCP transport mode",
    )
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> None:
    global _STATE

    args = parse_args(argv)

    if mcp is None:
        print(
            "Unable to start Aegis MCP server: mcp package import failed.",
            file=sys.stderr,
        )
        print(str(MCP_IMPORT_ERROR), file=sys.stderr)
        raise SystemExit(1)

    config_path = Path(args.config)
    schema_path = Path(args.schema) if args.schema else None

    try:
        context = load_and_validate_context(
            context_path=config_path,
            schema_path=schema_path,
        )
    except SecurityContextValidationError as exc:
        print(
            json.dumps(
                {
                    "error": "Invalid security-context.yaml",
                    "config_path": str(config_path),
                    "details": exc.errors,
                },
                indent=2,
            ),
            file=sys.stderr,
        )
        raise SystemExit(1)

    _STATE = AppState(config_path=config_path, schema_path=schema_path, context=context)

    mcp.run(transport=args.transport)


if __name__ == "__main__":
    main()
