#!/usr/bin/env python3
"""Merge partial OpenAPI YAML files into a single cohesive spec."""

import yaml
import sys
from pathlib import Path


def deep_merge(base: dict, overlay: dict) -> dict:
    """Recursively merge overlay into base, returning the merged result."""
    for key, value in overlay.items():
        if key in base and isinstance(base[key], dict) and isinstance(value, dict):
            deep_merge(base[key], value)
        else:
            base[key] = value
    return base


def resolve_local_refs(obj, source_file: str, all_schemas: dict):
    """
    Rewrite $ref pointers that reference schemas.yaml or local component schemas
    so they point to #/components/schemas/... in the merged spec.
    Also collect any local component schemas into all_schemas.
    """
    if isinstance(obj, dict):
        if "$ref" in obj:
            ref = obj["$ref"]
            # Rewrite references to schemas.yaml
            if ref.startswith("schemas.yaml#/components/schemas/"):
                schema_name = ref.split("/")[-1]
                obj["$ref"] = f"#/components/schemas/{schema_name}"
            elif ref.startswith("schemas.yaml#/components/parameters/"):
                param_name = ref.split("/")[-1]
                obj["$ref"] = f"#/components/parameters/{param_name}"
            elif ref.startswith("#/components/schemas/"):
                # Local schema ref — already correct format, but we need to
                # make sure the schema is collected
                pass
            return obj
        for key, value in obj.items():
            obj[key] = resolve_local_refs(value, source_file, all_schemas)
    elif isinstance(obj, list):
        for i, item in enumerate(obj):
            obj[i] = resolve_local_refs(item, source_file, all_schemas)
    return obj


def main():
    parts_dir = Path(__file__).parent / "api_parts"
    output_path = Path(__file__).parent / "duffel_openapi.yaml"

    # Load the schemas file first
    schemas_path = parts_dir / "schemas.yaml"
    with open(schemas_path) as f:
        schemas_doc = yaml.safe_load(f)

    # Build the base merged spec
    merged = {
        "openapi": "3.0.3",
        "info": {
            "title": "Duffel API",
            "description": (
                "The Duffel API provides a unified interface for searching, booking, "
                "and managing flights and accommodation. This spec covers the Flights, "
                "Stays, Payments, Notifications, Identity, and Supporting Resources APIs.\n\n"
                "**Base URLs:**\n"
                "- Main API: `https://api.duffel.com`\n"
                "- Card API (PCI compliant): `https://api.duffel.cards`\n\n"
                "**Authentication:** Bearer token via `Authorization` header.\n\n"
                "**Required headers:** `Duffel-Version: v2`, `Accept: application/json`.\n\n"
                "**Pagination:** Cursor-based with `limit`, `after`, `before` query parameters. "
                "Default 50, max 200 results per page.\n\n"
                "**Client timeout:** Set HTTP client timeout to at least 130 seconds "
                "for order/booking creation endpoints."
            ),
            "version": "v2",
            "contact": {
                "name": "Duffel",
                "url": "https://duffel.com",
            },
        },
        "servers": [
            {
                "url": "https://api.duffel.com",
                "description": "Duffel API (main)",
            },
            {
                "url": "https://api.duffel.cards",
                "description": "Duffel Cards API (PCI compliant, for card endpoints only)",
            },
        ],
        "security": [{"BearerAuth": []}],
        "paths": {},
        "components": {
            "securitySchemes": schemas_doc.get("components", {}).get("securitySchemes", {}),
            "parameters": schemas_doc.get("components", {}).get("parameters", {}),
            "schemas": schemas_doc.get("components", {}).get("schemas", {}),
        },
        "tags": [
            {"name": "Flights - Offer Requests", "description": "Search for flight offers."},
            {"name": "Flights - Offers", "description": "Browse and price flight offers."},
            {"name": "Flights - Orders", "description": "Book flights and manage orders."},
            {"name": "Flights - Payments", "description": "Pay for flight orders."},
            {"name": "Flights - Seat Maps", "description": "Retrieve seat maps for offers."},
            {"name": "Flights - Order Cancellations", "description": "Cancel flight orders."},
            {"name": "Flights - Order Changes", "description": "Change existing flight orders."},
            {"name": "Flights - Partial Offer Requests", "description": "Multi-step flight searches."},
            {"name": "Flights - Airline-Initiated Changes", "description": "Handle airline schedule changes."},
            {"name": "Flights - Batch Offer Requests", "description": "Async batch flight searches."},
            {"name": "Flights - Airline Credits", "description": "Manage airline credits."},
            {"name": "Stays - Search", "description": "Search for accommodation."},
            {"name": "Stays - Quotes", "description": "Get quotes for stays."},
            {"name": "Stays - Bookings", "description": "Book and manage stays."},
            {"name": "Stays - Accommodation", "description": "Browse accommodation data."},
            {"name": "Stays - Negotiated Rates", "description": "Manage negotiated rates."},
            {"name": "Stays - Brands", "description": "Browse accommodation brands."},
            {"name": "Stays - Loyalty Programmes", "description": "Accommodation loyalty programmes."},
            {"name": "Payments - Cards", "description": "Store payment cards (PCI-compliant host)."},
            {"name": "Payments - 3D Secure", "description": "3D Secure authentication sessions."},
            {"name": "Notifications - Webhooks", "description": "Manage webhook endpoints."},
            {"name": "Notifications - Webhook Deliveries", "description": "Track webhook deliveries."},
            {"name": "Notifications - Webhook Events", "description": "Browse webhook events."},
            {"name": "Identity - Customer Users", "description": "Manage customer users."},
            {"name": "Identity - Customer User Groups", "description": "Manage customer user groups."},
            {"name": "Identity - Component Client Keys", "description": "Generate UI component keys."},
            {"name": "Supporting Resources - Airlines", "description": "Airline reference data."},
            {"name": "Supporting Resources - Aircraft", "description": "Aircraft reference data."},
            {"name": "Supporting Resources - Airports", "description": "Airport reference data."},
            {"name": "Supporting Resources - Cities", "description": "City reference data."},
            {"name": "Supporting Resources - Places", "description": "Place search suggestions."},
            {"name": "Supporting Resources - Loyalty Programmes", "description": "Airline loyalty programmes."},
        ],
    }

    # List of endpoint part files to merge (order determines tag ordering in spec)
    part_files = [
        "flights_offer_requests.yaml",
        "flights_offers.yaml",
        "flights_orders.yaml",
        "flights_payments.yaml",
        "flights_seat_maps.yaml",
        "flights_order_cancellations.yaml",
        "flights_order_changes.yaml",
        "flights_partial_offer_requests.yaml",
        "flights_airline_initiated_changes.yaml",
        "flights_batch_offer_requests.yaml",
        "flights_airline_credits.yaml",
        "stays.yaml",
        "payments_and_3ds.yaml",
        "notifications.yaml",
        "identity.yaml",
        "supporting_resources.yaml",
    ]

    for filename in part_files:
        filepath = parts_dir / filename
        if not filepath.exists():
            print(f"WARNING: {filename} not found, skipping")
            continue

        with open(filepath) as f:
            doc = yaml.safe_load(f)

        if not doc:
            continue

        # Collect any local component schemas
        local_schemas = doc.get("components", {}).get("schemas", {})
        if local_schemas:
            # Resolve refs within local schemas first
            resolve_local_refs(local_schemas, filename, {})
            merged["components"]["schemas"].update(local_schemas)

        # Merge paths
        paths = doc.get("paths", {})
        if paths:
            resolve_local_refs(paths, filename, {})
            for path, methods in paths.items():
                if path in merged["paths"]:
                    deep_merge(merged["paths"][path], methods)
                else:
                    merged["paths"][path] = methods

    # Add common error responses to all operations
    error_responses = {
        "400": {
            "description": "Bad request.",
            "content": {
                "application/json": {
                    "schema": {"$ref": "#/components/schemas/ErrorResponse"}
                }
            },
        },
        "401": {
            "description": "Unauthorized.",
            "content": {
                "application/json": {
                    "schema": {"$ref": "#/components/schemas/ErrorResponse"}
                }
            },
        },
        "422": {
            "description": "Validation error.",
            "content": {
                "application/json": {
                    "schema": {"$ref": "#/components/schemas/ErrorResponse"}
                }
            },
        },
        "429": {
            "description": "Rate limit exceeded.",
            "content": {
                "application/json": {
                    "schema": {"$ref": "#/components/schemas/ErrorResponse"}
                }
            },
        },
    }

    for path, methods in merged["paths"].items():
        for method, operation in methods.items():
            if isinstance(operation, dict) and "responses" in operation:
                for code, resp in error_responses.items():
                    if code not in operation["responses"]:
                        operation["responses"][code] = resp

    # Write the merged spec
    # Use a custom representer to handle multiline strings nicely
    class LiteralStr(str):
        pass

    def literal_representer(dumper, data):
        if "\n" in data:
            return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
        return dumper.represent_scalar("tag:yaml.org,2002:str", data)

    yaml.add_representer(LiteralStr, literal_representer)

    with open(output_path, "w") as f:
        yaml.dump(
            merged,
            f,
            default_flow_style=False,
            sort_keys=False,
            allow_unicode=True,
            width=120,
        )

    # Count stats
    n_paths = len(merged["paths"])
    n_operations = sum(
        1
        for methods in merged["paths"].values()
        for method, op in methods.items()
        if isinstance(op, dict) and "operationId" in op
    )
    n_schemas = len(merged["components"]["schemas"])

    print(f"Merged OpenAPI spec written to: {output_path}")
    print(f"  Paths:      {n_paths}")
    print(f"  Operations: {n_operations}")
    print(f"  Schemas:    {n_schemas}")


if __name__ == "__main__":
    main()
