import json
import sys
from pathlib import Path
from typing import Any


ROOT_PATH = Path(__file__).resolve().parents[1]
OUTPUT_PATH = ROOT_PATH / "web" / "openapi.json"
FRONTEND_CONTRACT_CONSTANTS_PATH = ROOT_PATH / "web" / "src" / "shared" / "api" / "generated" / "constants.ts"

if str(ROOT_PATH) not in sys.path:
    sys.path.insert(0, str(ROOT_PATH))


def export_openapi_schema() -> Path:
    from app import create_app

    app = create_app()
    schema = app.openapi()
    _patch_auth_contracts(app, schema)
    _patch_validation_contracts(schema)
    _patch_error_contracts(schema)
    _register_extra_schemas(schema)
    OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    OUTPUT_PATH.write_text(
        json.dumps(schema, ensure_ascii=False, indent=2, sort_keys=True),
        encoding="utf-8",
    )
    export_frontend_contract_constants(schema)
    return OUTPUT_PATH


def export_frontend_contract_constants(schema: dict[str, Any]) -> Path:
    from schema.work_project.graph import EDGE_TYPE_CATEGORY, WorkProjectGraphEdgeCategory

    default_command = _get_schema_property_default(
        schema,
        "CreateSandboxContainerRequest",
        "container_command",
    )
    if not isinstance(default_command, str):
        raise RuntimeError("CreateSandboxContainerRequest.container_command default must be a string")

    # Edge category is derived from edge type in Python, so it is serialized from that single source.
    edge_categories = [category.value for category in WorkProjectGraphEdgeCategory]
    edge_type_category = {edge_type.value: category.value for edge_type, category in EDGE_TYPE_CATEGORY.items()}

    FRONTEND_CONTRACT_CONSTANTS_PATH.parent.mkdir(parents=True, exist_ok=True)
    FRONTEND_CONTRACT_CONSTANTS_PATH.write_text(
        "// This file was auto-generated by scripts/export_schema.py.\n"
        "// Do not make direct changes to this file.\n\n"
        f"export const SYSTEM_USER_ROLES = {_enum_values_ts(schema, 'SystemUserRole')} as const;\n"
        f"export const WORK_PROJECT_TYPES = {_enum_values_ts(schema, 'WorkProjectType')} as const;\n"
        f"export const WORK_PROJECT_STATUSES = {_enum_values_ts(schema, 'WorkProjectStatus')} as const;\n"
        f"export const WORK_PROJECT_STATUS = {_enum_object_ts(schema, 'WorkProjectStatus')} as const;\n"
        f"export const WORK_PROJECT_TASK_STATUSES = {_enum_values_ts(schema, 'WorkProjectTaskStatus')} as const;\n"
        f"export const WORK_PROJECT_ASSET_TYPES = {_enum_values_ts(schema, 'WorkProjectAssetType')} as const;\n"
        f"export const WORK_PROJECT_ASSET_TYPE = {_enum_object_ts(schema, 'WorkProjectAssetType')} as const;\n"
        f"export const WORK_PROJECT_ASSET_ORIGINS = {_enum_values_ts(schema, 'WorkProjectAssetOrigin')} as const;\n"
        f"export const WORK_PROJECT_ASSET_ORIGIN = {_enum_object_ts(schema, 'WorkProjectAssetOrigin')} as const;\n"
        f"export const WORK_PROJECT_FINDING_SEVERITIES = {_enum_values_ts(schema, 'WorkProjectFindingSeverity')} as const;\n"
        f"export const WORK_PROJECT_FINDING_STATUSES = {_enum_values_ts(schema, 'WorkProjectFindingStatus')} as const;\n"
        f"export const WORK_PROJECT_GRAPH_EDGE_TYPES = {_enum_values_ts(schema, 'WorkProjectGraphEdgeType')} as const;\n"
        f"export const WORK_PROJECT_GRAPH_EDGE_CATEGORIES = {json.dumps(edge_categories)} as const;\n"
        f"export const WORK_PROJECT_GRAPH_EDGE_CATEGORY = {json.dumps(edge_type_category)} as const;\n"
        f"export const WORK_PROJECT_ATTACK_PATH_STATUSES = {_enum_values_ts(schema, 'WorkProjectAttackPathStatus')} as const;\n"
        f"export const SANDBOX_IMAGE_STATUSES = {_enum_values_ts(schema, 'SandboxImageStatus')} as const;\n"
        f"export const SANDBOX_CONTAINER_STATUSES = {_enum_values_ts(schema, 'SandboxContainerStatus')} as const;\n"
        f"export const SESSION_TYPES = {_enum_values_ts(schema, 'SessionType')} as const;\n\n"
        f"export const ACCESS_TOKEN_HEADER = {json.dumps(_get_access_token_header())};\n"
        f"export const SANDBOX_CONTAINER_DEFAULT_COMMAND = {json.dumps(default_command)};\n",
        encoding="utf-8",
    )
    return FRONTEND_CONTRACT_CONSTANTS_PATH


def _enum_values_ts(schema: dict[str, Any], schema_name: str) -> str:
    values = schema.get("components", {}).get("schemas", {}).get(schema_name, {}).get("enum")
    if not isinstance(values, list) or not all(isinstance(value, str) for value in values):
        raise RuntimeError(f"OpenAPI schema is missing string enum {schema_name}")
    return json.dumps(values)


def _enum_object_ts(schema: dict[str, Any], schema_name: str) -> str:
    values = schema.get("components", {}).get("schemas", {}).get(schema_name, {}).get("enum")
    if not isinstance(values, list) or not all(isinstance(value, str) for value in values):
        raise RuntimeError(f"OpenAPI schema is missing string enum {schema_name}")
    return json.dumps({_enum_member_name(value): value for value in values})


def _enum_member_name(value: str) -> str:
    return value.upper().replace("-", "_")


def _get_schema_property_default(schema: dict[str, Any], schema_name: str, property_name: str) -> Any:
    try:
        return schema["components"]["schemas"][schema_name]["properties"][property_name]["default"]
    except KeyError as exc:
        raise RuntimeError(f"OpenAPI schema is missing {schema_name}.{property_name}.default") from exc


def _get_access_token_header() -> str:
    from middleware.auth import ACCESS_TOKEN_HEADER

    return ACCESS_TOKEN_HEADER


def _patch_auth_contracts(app: Any, schema: dict[str, Any]) -> None:
    """drive OpenAPI security from the FastAPI dependency tree.

    routes that depend on require_user get AccessTokenAuth + 401; routes that
    additionally depend on require_admin also get 403."""
    from middleware.auth import require_admin, require_user

    components = schema.setdefault("components", {})
    schemas = components.setdefault("schemas", {})
    schemas.setdefault("CommonResponse_Any_", _common_response_any_schema())
    components.setdefault("securitySchemes", {})["AccessTokenAuth"] = {
        "type": "apiKey",
        "in": "header",
        "name": _get_access_token_header(),
    }

    paths = schema.get("paths", {})
    for route in app.routes:
        path = getattr(route, "path", "")
        methods = getattr(route, "methods", None)
        if not methods or path not in paths:
            continue

        deps = _route_dependency_funcs(route)
        if not (require_user in deps or require_admin in deps):
            continue

        for method in methods:
            operation = paths[path].get(method.lower())
            if operation is None:
                continue
            operation["security"] = [{"AccessTokenAuth": []}]
            responses = operation.setdefault("responses", {})
            responses["401"] = _common_response_ref("Unauthorized")
            if require_admin in deps:
                responses["403"] = _common_response_ref("Forbidden")


def _route_dependency_funcs(route: Any) -> set:
    funcs: set = set()
    dependant = getattr(route, "dependant", None)
    if dependant is not None:
        _walk_dependant(dependant, funcs)
    return funcs


def _walk_dependant(dependant: Any, funcs: set) -> None:
    if dependant.call is not None:
        funcs.add(dependant.call)
    for sub in dependant.dependencies:
        _walk_dependant(sub, funcs)


def _patch_validation_contracts(schema: dict[str, Any]) -> None:
    """rewrite FastAPI's default 422 docs to the runtime CommonResponse shape"""
    for methods in schema.get("paths", {}).values():
        for operation in methods.values():
            responses = operation.get("responses")
            if isinstance(responses, dict) and "422" in responses:
                _ensure_common_response(responses["422"], "Validation Error")


def _patch_error_contracts(schema: dict[str, Any]) -> None:
    """normalize documented error responses to the runtime CommonResponse shape"""
    descriptions = {
        "400": "Bad Request",
        "401": "Unauthorized",
        "403": "Forbidden",
        "404": "Not Found",
        "422": "Validation Error",
    }
    for methods in schema.get("paths", {}).values():
        for operation in methods.values():
            responses = operation.get("responses")
            if not isinstance(responses, dict):
                continue
            for status_code, description in descriptions.items():
                response = responses.get(status_code)
                if isinstance(response, dict):
                    _ensure_common_response(response, description)


def _register_extra_schemas(schema: dict[str, Any]) -> None:
    """publish ws contracts as OpenAPI components so the frontend can derive types"""
    from pydantic import TypeAdapter
    from schema.agent.events import (
        AgentStreamCancelAllCommand,
        AgentStreamCommandSchema,
        AgentStreamInterruptCommand,
        AgentStreamSendCommand,
        AgentEventSchema,
        DoneEvent,
        RunStateEvent,
    )
    from schema.agent.subordinates import AgentSubordinateTaskToolItem, AgentSubordinateTaskToolResult
    from schema.common.tool_results import ToolResultSchema
    from schema.sandbox.command_outputs import (
        SandboxCommandOutputChunk,
        SandboxCommandResultMetadata,
    )

    components = schema.setdefault("components", {})
    schemas = components.setdefault("schemas", {})

    extras = {
        "DoneEvent": DoneEvent,
        "RunStateEvent": RunStateEvent,
        "AgentEventSchema": AgentEventSchema,
        "AgentStreamSendCommand": AgentStreamSendCommand,
        "AgentStreamInterruptCommand": AgentStreamInterruptCommand,
        "AgentStreamCancelAllCommand": AgentStreamCancelAllCommand,
        "AgentStreamCommandSchema": AgentStreamCommandSchema,
        "AgentSubordinateTaskToolItem": AgentSubordinateTaskToolItem,
        "AgentSubordinateTaskToolResult": AgentSubordinateTaskToolResult,
        "ToolResultSchema": ToolResultSchema,
        "SandboxCommandResultMetadata": SandboxCommandResultMetadata,
        "SandboxCommandOutputChunk": SandboxCommandOutputChunk,
    }
    for name, model in extras.items():
        body = TypeAdapter(model).json_schema(ref_template="#/components/schemas/{model}")
        for nested_name, nested_body in body.pop("$defs", {}).items():
            schemas.setdefault(nested_name, nested_body)
        schemas[name] = body


def _common_response_ref(description: str) -> dict[str, Any]:
    return {
        "description": description,
        "content": {
            "application/json": {
                "schema": {"$ref": "#/components/schemas/CommonResponse_Any_"},
            },
        },
    }


def _ensure_common_response(response: dict[str, Any], description: str) -> None:
    response["description"] = response.get("description") or description
    content = response.setdefault("content", {})
    json_content = content.setdefault("application/json", {})
    json_content["schema"] = {"$ref": "#/components/schemas/CommonResponse_Any_"}


def _common_response_any_schema() -> dict[str, Any]:
    return {
        "type": "object",
        "title": "CommonResponse[Any]",
        "properties": {
            "code": {"type": "integer", "title": "Code", "default": 200},
            "message": {"type": "string", "title": "Message", "default": "success"},
            "data": {"title": "Data"},
        },
    }


if __name__ == "__main__":
    output_path = export_openapi_schema()
    print(f"OpenAPI schema exported to {output_path}")
    print(f"Frontend contract constants exported to {FRONTEND_CONTRACT_CONSTANTS_PATH}")
