#!/usr/bin/env python
# Copyright Salient Predictions 2026

"""OpenAPI Schema Utilities.

This module provides utilities for downloading the Salient API OpenAPI schema,
extracting enum definitions and documentation, and generating Python type stubs
for use with static type checkers and IDE autocompletion.

Usage:
    poetry run python -m salientsdk.schema

This will download the latest schema and regenerate `_generated_types.py`.

To generate files in a custom directory:
    poetry run python -m salientsdk.schema --dest /path/to/output
"""

import argparse
import re
import subprocess
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import requests

# Default paths
DEFAULT_BASE_URL = "https://api-beta.salientpredictions.com"
SCHEMA_SPEC_PATH = "/v2/documentation/spec"
DEFAULT_TYPES_OUTPUT_PATH = Path(__file__).parent / "_generated_types.py"

# Line length limit for generated code (to be compatible with black/ruff)
MAX_LINE_LENGTH = 99

# Endpoints to exclude from type generation (not yet public or deprecated)
EXCLUDED_ENDPOINTS = {"risk_summary", "risk", "reliability", "forecast_distribution"}


def download_schema(base_url: str = DEFAULT_BASE_URL) -> dict[str, Any]:
    """Download the OpenAPI schema from the Salient API.

    Fetches the OpenAPI 3.x specification from the API documentation endpoint.
    The full URL is constructed as: `{base_url}/v2/documentation/spec`

    Args:
        base_url: The base URL of the Salient API.
            Defaults to "https://api-beta.salientpredictions.com".
            Production would be "https://api.salientpredictions.com".

    Returns:
        The parsed OpenAPI schema as a dictionary.

    Raises:
        requests.HTTPError: If the request fails.
        json.JSONDecodeError: If the response is not valid JSON.

    Example:
        >>> schema = download_schema()
        >>> schema["info"]["title"]
        'Salient Predictions'
    """
    url = base_url.rstrip("/") + SCHEMA_SPEC_PATH
    response = requests.get(url, timeout=30)
    response.raise_for_status()
    return response.json()


def _parse_enum_descriptions(description: str) -> dict[str, str]:
    """Parse individual enum value descriptions from a parameter description.

    The OpenAPI schema often includes descriptions for each enum value in the format:
    - **value**: Description text here.

    Args:
        description: The full parameter description string.

    Returns:
        A dictionary mapping enum values to their descriptions.
    """
    enum_docs = {}
    # Match patterns like "- **value**: description text"
    # The description continues until the next "- **" or end of string
    pattern = r"-\s*\*\*([^*]+)\*\*:\s*(.+?)(?=\n-\s*\*\*|$)"
    matches = re.findall(pattern, description, re.DOTALL)
    for value, desc in matches:
        # Clean up the description - collapse whitespace
        desc = re.sub(r"\s+", " ", desc).strip()
        if desc:
            enum_docs[value.strip()] = desc
    return enum_docs


def extract_endpoint_metadata(schema: dict[str, Any]) -> dict[str, dict[str, Any]]:
    """Extract endpoint metadata including enums and documentation from the schema.

    Parses the OpenAPI schema and extracts, for each endpoint:
    - The endpoint's description (function docstring)
    - Each parameter's enum values (if any)
    - Each parameter's description (enum documentation)

    The returned structure is organized by endpoint name (e.g., "forecast_timeseries")
    and contains all information needed to generate type stubs.

    Args:
        schema: The OpenAPI schema dictionary.

    Returns:
        A dictionary keyed by endpoint name (derived from operationId), where each value
        contains:
        - "description": The endpoint's description string
        - "parameters": A dict of parameter names to their metadata:
            - "enum": List of valid enum values (if applicable)
            - "description": The parameter's description string
            - "type": The parameter's type (string, array, etc.)

    Example:
        >>> schema = download_schema()
        >>> metadata = extract_endpoint_metadata(schema)
        >>> metadata["forecast_timeseries"]["description"][:50]
        'Return a timeseries of forecast probability quanti'
        >>> metadata["forecast_timeseries"]["parameters"]["model"]["enum"][:3]
        ['ecmwf_seas5_calib', 'gfs', 'noaa_gefs_calib']
        >>> "blend" in metadata["forecast_timeseries"]["parameters"]["model"]["description"]
        True
    """
    metadata = {}
    paths = schema.get("paths", {})

    for path, path_item in paths.items():
        # Process each HTTP method (get, post, etc.)
        for method, operation in path_item.items():
            if method in ("parameters", "servers", "summary", "description"):
                # Skip non-method keys
                continue
            if not isinstance(operation, dict):
                continue

            # Get operation ID (e.g., "forecast_timeseries", "get_data_timeseries")
            operation_id = operation.get("operationId")
            if not operation_id:
                continue

            # Normalize operation ID - remove common prefixes like "get_"
            endpoint_name = operation_id
            if endpoint_name.startswith("get_"):
                endpoint_name = endpoint_name[4:]

            endpoint_metadata = {
                "description": operation.get("description", ""),
                "parameters": {},
            }

            # Process parameters
            parameters = operation.get("parameters", [])
            for param in parameters:
                param_name = param.get("name")
                if not param_name:
                    continue

                param_schema = param.get("schema", {})
                param_description = param.get("description", "")

                # Handle array types with items
                if param_schema.get("type") == "array":
                    items = param_schema.get("items", {})
                    enum_values = items.get("enum")
                    param_type = "array"
                else:
                    enum_values = param_schema.get("enum")
                    param_type = param_schema.get("type", "string")

                # Only include parameters that have enum values
                if enum_values:
                    # Parse individual enum descriptions from the parameter description
                    enum_docs = _parse_enum_descriptions(param_description)

                    endpoint_metadata["parameters"][param_name] = {
                        "enum": enum_values,
                        "description": param_description,
                        "enum_docs": enum_docs,
                        "type": param_type,
                    }

            # Only include endpoints that have at least one enum parameter
            if endpoint_metadata["parameters"]:
                metadata[endpoint_name] = endpoint_metadata

    return metadata


def _to_pascal_case(snake_str: str) -> str:
    """Convert a snake_case or kebab-case string to PascalCase.

    Args:
        snake_str: The input string in snake_case or kebab-case.

    Returns:
        The string converted to PascalCase.

    Examples:
        >>> _to_pascal_case("forecast_timeseries")
        'ForecastTimeseries'
        >>> _to_pascal_case("sub-seasonal")
        'SubSeasonal'
    """
    # Replace hyphens with underscores, then split and capitalize
    parts = snake_str.replace("-", "_").split("_")
    return "".join(word.capitalize() for word in parts)


def _format_literal_value(v: Any) -> str:
    """Format a value for use in a Literal type annotation.

    Args:
        v: The value to format.

    Returns:
        The formatted string representation.
    """
    if isinstance(v, str):
        return f'"{v}"'
    else:
        return str(v)


def _format_literal_type(type_name: str, enum_values: list[Any]) -> list[str]:
    """Format a Literal type declaration, breaking lines if necessary.

    Args:
        type_name: The name of the type alias.
        enum_values: The list of enum values.

    Returns:
        A list of lines for the type declaration.
    """
    formatted_values = [_format_literal_value(v) for v in enum_values]

    # Try single line first
    single_line = f"{type_name} = Literal[{', '.join(formatted_values)}]"
    if len(single_line) <= MAX_LINE_LENGTH:
        return [single_line]

    # Multi-line format
    lines = [f"{type_name} = Literal["]
    for i, value in enumerate(formatted_values):
        comma = "," if i < len(formatted_values) - 1 else ","
        lines.append(f"    {value}{comma}")
    lines.append("]")
    return lines


def generate_types_module(
    metadata: dict[str, dict[str, Any]],
    output_path: Path = DEFAULT_TYPES_OUTPUT_PATH,
) -> Path:
    """Generate a Python module containing Literal types with docstrings.

    Creates a `_generated_types.py` file with:

    1. **Section headers** for each endpoint including the endpoint description:
       ```
       # ======================================================================
       # forecast_timeseries
       # Return a timeseries of forecast probability quantiles or ensembles.
       # ======================================================================
       ```

    2. **Literal type aliases** for each endpoint's enum parameters:
       - Named as `{EndpointName}{ParameterName}` in PascalCase
       - Example: `ForecastTimeseriesModel = Literal["gem", "clim", "blend", ...]`
       - Each Literal is followed by a docstring containing the full parameter
         description from the API schema.

    The generated file includes a header warning that it is auto-generated
    and should not be edited manually.

    Args:
        metadata: The endpoint metadata as returned by `extract_endpoint_metadata()`.
        output_path: The file path to write the generated module to.
            Defaults to `salientsdk/_generated_types.py`.

    Returns:
        The path where the module was written.

    Example:
        >>> schema = download_schema()
        >>> metadata = extract_endpoint_metadata(schema)
        >>> path = generate_types_module(metadata)
        >>> print(f"Generated {path}")
    """
    output_path = Path(output_path)
    now = datetime.now(timezone.utc)
    year = now.strftime("%Y")
    timestamp = now.strftime("%Y-%m-%dT%H:%M:%SZ")

    lines = [
        "#!/usr/bin/env python",
        f"# Copyright Salient Predictions {year}",
        "",
        '"""Auto-generated type definitions from the Salient API OpenAPI schema.',
        "",
        "DO NOT EDIT THIS FILE MANUALLY.",
        "",
        "Regenerate with: poetry run python -m salientsdk.schema",
        f"Generated: {timestamp}",
        '"""',
        "",
        "from typing import Literal",
        "",
    ]

    # Track all generated variable names for the get_all_types function
    generated_names: list[str] = []

    # Sort endpoints for consistent output
    for endpoint_name in sorted(metadata.keys()):
        endpoint_data = metadata[endpoint_name]
        endpoint_pascal = _to_pascal_case(endpoint_name)

        # Header with endpoint description
        endpoint_desc = endpoint_data.get("description", "")
        endpoint_first_line = endpoint_desc.split("\n")[0].strip() if endpoint_desc else ""
        lines.append(f"# {'=' * 70}")
        lines.append(f"# {endpoint_name}")
        if endpoint_first_line:
            lines.append(f"# {endpoint_first_line}")
        lines.append(f"# {'=' * 70}")
        lines.append("")

        # Process each parameter with enums
        for param_name in sorted(endpoint_data["parameters"].keys()):
            param_data = endpoint_data["parameters"][param_name]
            param_pascal = _to_pascal_case(param_name)
            type_name = f"{endpoint_pascal}{param_pascal}"

            enum_values = sorted(param_data["enum"], key=str)  # Sort for deterministic output

            # Generate Literal type (with line breaks if needed)
            generated_names.append(type_name)
            literal_lines = _format_literal_type(type_name, enum_values)
            lines.extend(literal_lines)

            # Generate docstring for the Literal
            param_desc = param_data.get("description", "")
            # Normalize whitespace in description
            param_desc = re.sub(r"[ \t]+", " ", param_desc)  # Collapse horizontal whitespace
            param_desc = re.sub(r"\n ", "\n", param_desc)  # Remove leading space after newlines
            # Strip trailing whitespace from each line (for linter compliance)
            param_desc = "\n".join(line.rstrip() for line in param_desc.split("\n"))
            param_desc = param_desc.strip()
            # Convert **value**: to `value`: (backticks render in IDEs, bold doesn't)
            param_desc = re.sub(r"\*\*([^*]+)\*\*:", r"`\1`:", param_desc)
            # Escape any triple quotes in the description
            param_desc = param_desc.replace('"""', "'''")

            lines.append(f'"""{endpoint_name} {param_name} options.')
            if param_desc:
                lines.append("")
                lines.append(param_desc)
            lines.append('"""')
            lines.append("")

    # Generate a function that returns all types (helps with vulture unused variable warnings)
    lines.append("")
    lines.append("def get_all_types() -> dict[str, object]:")
    lines.append('    """Return all generated Literal types as a dictionary.')
    lines.append("")
    lines.append("    This function exists to satisfy vulture's unused variable detection.")
    lines.append("    It also provides a convenient way to introspect all available types.")
    lines.append("")
    lines.append("    Returns:")
    lines.append("        A dictionary mapping Literal type names to their type aliases.")
    lines.append('    """')
    lines.append("    return {")
    for name in generated_names:
        lines.append(f'        "{name}": {name},')
    lines.append("    }")
    lines.append("")

    # Write the file (strip trailing whitespace for linter compliance)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w", encoding="utf-8") as f:
        f.write("\n".join(line.rstrip() for line in lines))

    # Format with black if available (for consistent output)
    try:
        subprocess.run(
            ["black", "--quiet", str(output_path)],
            check=False,
            capture_output=True,
        )
    except FileNotFoundError:
        pass  # black not installed, skip formatting

    return output_path


# Default path for generated docs
DEFAULT_DOCS_OUTPUT_PATH = Path(__file__).parent.parent / "docs" / "types.md"


def generate_types_docs(
    metadata: dict[str, dict[str, Any]],
    output_path: Path = DEFAULT_DOCS_OUTPUT_PATH,
) -> Path:
    """Generate a markdown documentation file for the generated types.

    Creates a `types.md` file organized by endpoint, suitable for mkdocs/mkdocstrings.

    Args:
        metadata: The endpoint metadata as returned by `extract_endpoint_metadata()`.
        output_path: The file path to write the markdown to.
            Defaults to `docs/types.md`.

    Returns:
        The path where the markdown was written.
    """
    output_path = Path(output_path)

    lines = [
        "# Enumerated Type Definitions",
        "",
        "These `Literal` lists describe the available options to API functions.",
        "",
    ]

    # Sort endpoints for consistent output
    for endpoint_name in sorted(metadata.keys()):
        endpoint_data = metadata[endpoint_name]
        endpoint_pascal = _to_pascal_case(endpoint_name)

        # Section header with link back to API docs
        lines.append(f"## [{endpoint_name}](api.md#salientsdk.{endpoint_name})")
        lines.append("")

        # Add each parameter type
        for param_name in sorted(endpoint_data["parameters"].keys()):
            param_pascal = _to_pascal_case(param_name)
            type_name = f"{endpoint_pascal}{param_pascal}"
            lines.append(f"::: salientsdk._generated_types.{type_name}")
            lines.append("")

    # Write the file
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))

    return output_path


def sync_schema(
    base_url: str = DEFAULT_BASE_URL,
    dest: Path | None = None,
    verbose: bool = True,
) -> Path:
    """Download the schema and regenerate the types module and docs.

    This is the main entry point that orchestrates the full sync process:
    1. Downloads the schema from the API
    2. Extracts endpoint metadata
    3. Generates the `_generated_types.py` module
    4. Generates the `docs/types.md` documentation

    Args:
        base_url: The base URL of the Salient API.
            Defaults to "https://api-beta.salientpredictions.com".
        dest: Destination directory for the generated files.
            If None, uses the default locations (salientsdk/ and docs/).
        verbose: If True, print progress messages.

    Returns:
        The path to the generated types module.

    Example:
        Run with poetry to sync the schema::

            poetry run python -m salientsdk.schema

        Or specify a custom destination::

            poetry run python -m salientsdk.schema --dest /tmp/output
    """
    types_output_path = Path(dest) / "_generated_types.py" if dest else DEFAULT_TYPES_OUTPUT_PATH
    docs_output_path = Path(dest) / "types.md" if dest else DEFAULT_DOCS_OUTPUT_PATH

    url = base_url.rstrip("/") + SCHEMA_SPEC_PATH

    if verbose:
        print(f"Downloading schema from {url}...")

    schema = download_schema(base_url)

    metadata = extract_endpoint_metadata(schema)

    # Filter out excluded endpoints
    metadata = {k: v for k, v in metadata.items() if k not in EXCLUDED_ENDPOINTS}

    if verbose:
        print(f"Extracted metadata for {len(metadata)} endpoints")

    output_path = generate_types_module(metadata, types_output_path)

    if verbose:
        print(f"Generated {output_path}")

    docs_path = generate_types_docs(metadata, docs_output_path)

    if verbose:
        print(f"Generated {docs_path}")

    return output_path


def main() -> None:
    """Command-line entry point for schema sync.

    Usage:
        poetry run python -m salientsdk.schema
        poetry run python -m salientsdk.schema --dest /path/to/output
        poetry run python -m salientsdk.schema --url https://api.salientpredictions.com
    """
    parser = argparse.ArgumentParser(
        description="Download Salient API schema and generate type definitions.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  poetry run python -m salientsdk.schema
  poetry run python -m salientsdk.schema --dest /tmp/output
  poetry run python -m salientsdk.schema --url https://api.salientpredictions.com
        """,
    )
    parser.add_argument(
        "--url",
        default=DEFAULT_BASE_URL,
        help=f"Base URL of the Salient API (default: {DEFAULT_BASE_URL})",
    )
    parser.add_argument(
        "--dest",
        type=Path,
        default=None,
        help="Destination directory for _generated_types.py (default: salientsdk package dir)",
    )
    parser.add_argument(
        "--quiet",
        "-q",
        action="store_true",
        help="Suppress progress messages",
    )

    args = parser.parse_args()

    sync_schema(
        base_url=args.url,
        dest=args.dest,
        verbose=not args.quiet,
    )


if __name__ == "__main__":
    main()
