engine marzano(0.1)
language python

pattern rename_resource() {
    or {
        `Audio` => `audio`,
        `ChatCompletion` => `chat.completions`,
        `Completion` => `completions`,
        `Edit` => `edits`,
        `Embedding` => `embeddings`,
        `File` => `files`,
        `FineTune` => `fine_tunes`,
        `FineTuningJob` => `fine_tuning`,
        `Image` => `images`,
        `Model` => `models`,
        `Moderation` => `moderations`,
    }
}

pattern rename_resource_cls() {
    or {
        r"Audio" => `resources.Audio`,
        r"ChatCompletion" => `resources.chat.Completions`,
        r"Completion" => `resources.Completions`,
        r"Edit" => `resources.Edits`,
        r"Embedding" => `resources.Embeddings`,
        r"File" => `resources.Files`,
        r"FineTune" => `resources.FineTunes`,
        r"FineTuningJob" => `resources.FineTuning`,
        r"Image" => `resources.Images`,
        r"Model" => `resources.Models`,
        r"Moderation" => `resources.Moderations`,
    }
}

pattern deprecated_resource() {
    or {
        `Customer`,
        `Deployment`,
        `Engine`,
        `ErrorObject`,
    }
}

pattern deprecated_resource_cls() {
    or {
        r"Customer",
        r"Deployment",
        r"Engine",
        r"ErrorObject",
    }
}


pattern rename_func($has_sync, $has_async, $res, $stmt, $params) {
    $func where {
        if ($func <: r"a([a-zA-Z0-9]+)"($func_rest)) {
            $has_async = `true`,
            $func => $func_rest,
            $stmt => `aclient.$res.$func($params)`,
        } else {
            $has_sync = `true`,
            $stmt => `client.$res.$func($params)`,
        },
        // Fix function renames
        if ($res <: `Image`) {
          $func => `generate`
        }
    }
}

pattern change_import($has_sync, $has_async, $need_openai_import) {
    $stmt where {
        $imports_and_defs = [],

        if ($need_openai_import <:  `true`) {
            $imports_and_defs += `import openai`,
        },

        if (and { $has_sync <: `true`, $has_async <: `true` }) {
            $imports_and_defs += `from openai import OpenAI, AsyncOpenAI`,
            $imports_and_defs += ``, // Blank line
            $imports_and_defs += `client = OpenAI()`,
            $imports_and_defs += `aclient = AsyncOpenAI()`,
        } else if ($has_sync <: `true`) {
            $imports_and_defs += `from openai import OpenAI`,
            $imports_and_defs += ``, // Blank line
            $imports_and_defs += `client = OpenAI()`,
        } else if ($has_async <: `true`) {
            $imports_and_defs += `from openai import AsyncOpenAI`,
            $imports_and_defs += ``, // Blank line
            $imports_and_defs += `aclient = AsyncOpenAI()`,
        },

        $separator = `\n`,
        $formatted = join(list = $imports_and_defs, $separator),
        $stmt => `$formatted`,
    }
}

pattern rewrite_whole_fn_call($import, $has_sync, $has_async, $res, $func, $params, $stmt, $body) {
    or {
        rename_resource() where {
            $import = `true`,
            $func <: rename_func($has_sync, $has_async, $res, $stmt, $params),
        },
        deprecated_resource() as $dep_res where {
            $stmt_whole = $stmt,
            if ($body <: contains `$_ = $stmt` as $line) {
                $stmt_whole = $line,
            },
            $stmt_whole => `# TODO: The resource '$dep_res' has been deprecated\n$stmt_whole`,
        }
    }
}

pattern unittest_patch() {
    or {
        decorated_definition($decorators, definition=$_) where {
            $decorators <: contains bubble decorator(value=`patch($cls_path)`) as $stmt where {
                $cls_path <: contains r"openai\.([a-zA-Z0-9]+)(?:.[^,]+)?"($res),
                if ($res <: rename_resource_cls()) {} else {
                    $res <: deprecated_resource_cls(),
                    $stmt => `# TODO: The resource '$res' has been deprecated\n$stmt`,
                }
            }
        },
        function_definition($body) where {
            $body <: contains bubble($body) or {
                `patch.object($params)`,
                `patch($params)`,
            } as $stmt where {
                $params <: contains bubble($body, $stmt) r"openai\.([a-zA-Z0-9]+)(?:.[^,]+)?"($res) where or {
                    $res <: rename_resource_cls(),
                    and {
                        $res <: deprecated_resource_cls(),
                        $line = $stmt,
                        if ($body <: contains or { `with $stmt:`, `with $stmt as $_:` } as $l) {
                            $line = $l,
                        },
                        $line => `# TODO: The resource '$res' has been deprecated\n$line`,
                    }
                }
            },
        }
    }
}

pattern pytest_patch() {
    decorated_definition($decorators, $definition) where {
        $decorators <: contains decorator(value=`pytest.fixture`),
        $definition <: bubble function_definition($body, $parameters) where {
            $parameters <: [$monkeypatch, ...],
            $body <: contains bubble($monkeypatch) or {
                `$monkeypatch.setattr($params)` as $stmt where {
                    $params <: contains bubble($stmt) r"openai\.([a-zA-Z0-9]+)(?:.[^,]+)?"($res) where or {
                        $res <: rename_resource_cls(),
                        $stmt => `# TODO: The resource '$res' has been deprecated\n$stmt`,
                    }
                },
                `monkeypatch.delattr($params)` as $stmt where {
                    $params <: contains bubble($stmt) r"openai\.([a-zA-Z0-9]+)(?:.[^,]+)?"($res) where or {
                        $res <: rename_resource_cls(),
                        $stmt => `# TODO: The resource '$res' has been deprecated\n$stmt`,
                    }
                },
            }
        },
    },
}

file($body) where {
    $need_openai_import = `false`,
    $has_openai_import = `false`,
    $has_partial_import = `false`,
    $has_sync = `false`,
    $has_async = `false`,

    // Remap errors
    $body <: maybe contains `openai.error.$exp` => `openai.$exp` where {
        $need_openai_import = `true`,
    },

    // Mark all the places where we they configure openai as something that requires manual intervention
    $body <: maybe contains bubble($need_openai_import) `openai.$field = $val` => `raise Exception("The 'openai.$field' option isn't read in the client API. You will need to pass it when you instantiate the client, e.g. 'OpenAI($field=$val)'")` where {
        $need_openai_import = `true`,
    },

    $body <: maybe contains `import openai` as $import_stmt where {
        $body <: contains bubble($has_sync, $has_async, $has_openai_import, $body) `openai.$res.$func($params)` as $stmt where {
            $res <: rewrite_whole_fn_call(import = $has_openai_import, $has_sync, $has_async, $res, $func, $params, $stmt, $body),
        },
    },

    $body <: maybe contains `from openai import $resources` as $partial_import_stmt where {
        $has_partial_import = `true`,
        $body <: contains bubble($has_sync, $has_async, $resources) `$res.$func($params)` as $stmt where {
            $resources <: contains $res,
            $res <: rewrite_whole_fn_call($import, $has_sync, $has_async, $res, $func, $params, $stmt, $body),
        }
    },

    if ($has_openai_import <: `true`) {
        $import_stmt <: change_import($has_sync, $has_async, $need_openai_import),
        if ($has_partial_import <: `true`) {
            $partial_import_stmt => .,
        },
    } else if ($has_partial_import <: `true`) {
        $partial_import_stmt <: change_import($has_sync, $has_async, $need_openai_import),
    },

    $body <: maybe contains unittest_patch(),
    $body <: maybe contains pytest_patch(),
}