# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import inspect
import os
from collections import OrderedDict
from contextlib import contextmanager
from dataclasses import dataclass
from typing import TYPE_CHECKING

import paddle
from paddle.incubate.cc.tools import apy_to_axpr_json
from paddle.static import InputSpec

from . import typing as pct

if TYPE_CHECKING:
    from collections.abc import Callable

__all__ = ['compile']


# Usage:
# import paddle.incubate.cc.typing as pct
# import paddle.incubate.cc as pcc
# import paddle.nn.functional as F
#
# N = pct.DimVar('N', min=2)
# K = pct.DimVar("K", min=2)
# M = pct.DimVar("M", 7168)
# DType = pct.DTypeVar("T", "bfloat16", "float32")
#
# def foo(
#   x: pct.Tensor([N, K], DType),
#   y: pct.Tensor([K, M], DType),
#   b: pct.Tensor([M], DType)
# ):
#   @pcc.force_register_fusion
#   def activate(out):
#     return F.relu(out + b)
#   return activate(x @ y)
#
# fused_foo = pcc.compile(
#   foo
# )
def compile(func, *args, **kwargs):
    annotations = _get_input_annotations(func)
    dtypes2func = {}
    for input_specs in _get_input_spec_lists(annotations):
        dtypes = tuple(input_spec.dtype for input_spec in input_specs)
        dtypes2func[dtypes] = _compile(func, input_specs, *args, **kwargs)
    return OverloadedFunc(FuncOverloadCtx(dtypes2func))


def _compile(
    func,
    input_specs,
    train=False,
    ap_path="",
    ap_workspace_dir='/tmp/paddle/ap',
    backend_device='cuda',
    target_framework='paddle',
    compile_engine='PCC',
):
    assert ap_path is not None
    assert not train, "only support inference now"
    assert backend_device in ["cuda", "dcu", "custom_device"]
    os.makedirs(ap_workspace_dir, exist_ok=True)
    build_strategy = paddle.static.BuildStrategy()
    assert compile_engine in ('CINN', 'PCC')
    with _ap_envs(ap_path, ap_workspace_dir, backend_device):
        static_fn = paddle.jit.to_static(
            func,
            input_spec=input_specs,
            build_strategy=build_strategy,
            full_graph=True,
            backend=compile_engine,
        )
        if not train:
            static_fn.eval()
        else:
            static_fn.train()
        concrete_program, partial_program_layer = (
            static_fn.get_concrete_program(
                *input_specs, is_train=static_fn._is_train_mode()
            )
        )
        partial_program_layer.training = static_fn._is_train_mode()
        # Force to generate the program immediately.
        if train:
            _ = partial_program_layer.train_program.forward_program
        else:
            _ = partial_program_layer.infer_program.forward_program
        return partial_program_layer


@dataclass
class FuncOverloadCtx:
    dtypes2func: dict[list[paddle.dtype], Callable]


class OverloadedFunc:
    def __init__(self, func_overload_ctx: FuncOverloadCtx):
        self.func_overload_ctx = func_overload_ctx

    def __call__(self, *args):
        dtypes = tuple(tensor.dtype for tensor in args)
        func = self.func_overload_ctx.dtypes2func.get(dtypes, None)
        assert func is not None, self.mismatched_debug_info(dtypes)
        return func(inputs=[*args])

    def mismatched_debug_info(self, dtypes):
        valid_signatures = "; ".join(
            f"[{idx + 1}] {dtypes}"
            for idx, pair in enumerate(
                self.func_overload_ctx.dtypes2func.items()
            )
            for dtypes in [pair[0]]
        )
        return f"input signature {dtypes} mismatched, valid signatures are: {valid_signatures}"


@dataclass
class InputSpecMakeCtx:
    name2dtype_num_candidates: dict[str, int]
    name2dtype_candidate_idx: dict[str, int]


@contextmanager
def _ap_envs(ap_path, ap_workspace_dir, backend_device):
    old_ap_workspace_dir = os.environ.get('AP_WORKSPACE_DIR')
    new_ap_path, old_ap_path = _get_ap_path(ap_path, backend_device)
    _convert_apy_to_axpr(new_ap_path)
    os.environ['AP_PATH'] = new_ap_path
    os.environ['AP_WORKSPACE_DIR'] = ap_workspace_dir
    new_flags, old_flags = _get_ap_flags()
    paddle.set_flags(new_flags)
    old_prim_all = paddle.base.core._is_all_prim_enabled()
    paddle.base.core._set_prim_all_enabled(True)
    try:
        yield
    finally:
        os.environ['AP_PATH'] = old_ap_path
        if old_ap_workspace_dir is not None:
            os.environ['AP_WORKSPACE_DIR'] = old_ap_workspace_dir
        paddle.set_flags(old_flags)
        paddle.base.core._set_prim_all_enabled(old_prim_all)


def _get_ap_path(ap_path, backend_device):
    ap_sys_path = f"{os.path.dirname(paddle.__file__)}/apy/sys"
    matmul_path = f"{os.path.dirname(paddle.__file__)}/apy/matmul_pass"
    if backend_device in ["cuda", "dcu"]:
        device_path = (
            f"{os.path.dirname(paddle.__file__)}/apy/device/{backend_device}"
        )
    else:
        device_path = ""
    old_ap_path = os.environ.get('AP_PATH')
    new_ap_path = f"{ap_sys_path}:{ap_path}:{device_path}:{matmul_path}:{old_ap_path if old_ap_path is not None else ''}"
    if old_ap_path is None:
        # Always add sys_path to AP_PATH, as it is required at runtime.
        old_ap_path = ap_sys_path
    return new_ap_path, old_ap_path


def _get_ap_flags():
    old_flags = paddle.get_flags(
        ['FLAGS_enable_ap', 'FLAGS_prim_enable_dynamic']
    )
    new_flags = dict(old_flags)
    new_flags['FLAGS_enable_ap'] = True
    new_flags['FLAGS_prim_enable_dynamic'] = True
    return new_flags, old_flags


def _convert_apy_to_axpr(ap_path):
    all_ap_paths = {p for p in ap_path.split(":") if p and os.path.isdir(p)}
    for path in all_ap_paths:
        apy_to_axpr_json.PyToAxpr(path)(path)


def _get_input_annotations(func):
    full_arg_spec = inspect.getfullargspec(func)
    return [
        pct_type
        for arg_name in full_arg_spec.args
        for pct_type in [full_arg_spec.annotations[arg_name]]
    ]


def _get_input_spec_lists(annotations):
    ctx = _create_empty_input_spec_make_ctx(annotations)
    assert len(ctx.name2dtype_num_candidates) > 0
    dtype_var_names = [
        pair[0] for pair in ctx.name2dtype_num_candidates.items()
    ]
    dtype_num_candidates = [
        pair[1] for pair in ctx.name2dtype_num_candidates.items()
    ]
    dtype_candidate_idx_compositions = _cartesian_product(
        [range(num_candidates) for num_candidates in dtype_num_candidates]
    )
    for idx_composition in dtype_candidate_idx_compositions:
        for arg_idx, candidate_idx in enumerate(idx_composition):
            ctx.name2dtype_candidate_idx[dtype_var_names[arg_idx]] = (
                candidate_idx
            )
        yield _get_input_specs(annotations, ctx)


def _create_empty_input_spec_make_ctx(annotations):
    ctx = InputSpecMakeCtx(OrderedDict(), OrderedDict())
    _init_empty_input_spec_make_ctx(annotations, ctx)
    return ctx


def _init_empty_input_spec_make_ctx(annotations, mut_ctx: InputSpecMakeCtx):
    for pct_type in annotations:
        _init_input_spec_make_ctx_name2dtype_num_candidates(pct_type, mut_ctx)


def _init_input_spec_make_ctx_name2dtype_num_candidates(
    pct_type, mut_ctx: InputSpecMakeCtx
):
    assert isinstance(pct_type.dtype, pct.DTypeVar), (
        f"pct_type.dtype should be a DTypeVar, but {type(pct_type.dtype)} were given."
    )
    name = pct_type.dtype.name
    if name in mut_ctx.name2dtype_num_candidates:
        assert mut_ctx.name2dtype_num_candidates[name] == len(
            pct_type.dtype.candidates
        )
    else:
        mut_ctx.name2dtype_num_candidates[name] = len(pct_type.dtype.candidates)


def _get_input_specs(annotations, ctx: InputSpecMakeCtx):
    return [_get_input_spec(pct_type, ctx) for pct_type in annotations]


def _get_input_spec(pct_type, ctx: InputSpecMakeCtx):
    assert isinstance(pct_type, pct.Tensor)
    return InputSpec(
        shape=_get_input_spec_shape(pct_type, ctx),
        dtype=_get_input_spec_dtype(pct_type, ctx),
    )


def _get_input_spec_shape(pct_type, ctx: InputSpecMakeCtx):
    return [_get_input_spec_shape_dim(dim_var) for dim_var in pct_type.shape]


def _get_input_spec_shape_dim(dim_var: pct.DimVar):
    if isinstance(dim_var, int):
        return dim_var
    assert isinstance(dim_var, pct.DimVar)
    if isinstance(dim_var.name_or_value, int):
        return dim_var.name_or_value
    return None


def _get_input_spec_dtype(pct_type, ctx: InputSpecMakeCtx):
    assert isinstance(pct_type.dtype, pct.DTypeVar)
    name = pct_type.dtype.name
    candidate_idx = ctx.name2dtype_candidate_idx[name]
    return pct_type.dtype.candidates[candidate_idx]


def _cartesian_product(lst_of_lst):
    assert len(lst_of_lst) > 0
    return _cartesian_product_impl([()], lst_of_lst)


def _cartesian_product_impl(collect_lst, lst_of_lst):
    if len(lst_of_lst) == 0:
        return collect_lst
    collect_lst = [(*x, y) for x in collect_lst for y in lst_of_lst[0]]
    return _cartesian_product_impl(collect_lst, lst_of_lst[1:])