import torch
from modules import config
from modules import generate_audio as generate

from functools import lru_cache
from typing import Callable

from modules.api.Api import APIManager

from modules.api.impl import (
    base_api,
    tts_api,
    ssml_api,
    google_api,
    openai_api,
    refiner_api,
)

torch._dynamo.config.cache_size_limit = 64
torch._dynamo.config.suppress_errors = True
torch.set_float32_matmul_precision("high")


def create_api():
    api = APIManager()

    base_api.setup(api)
    tts_api.setup(api)
    ssml_api.setup(api)
    google_api.setup(api)
    openai_api.setup(api)
    refiner_api.setup(api)

    return api


def conditional_cache(condition: Callable):
    def decorator(func):
        @lru_cache(None)
        def cached_func(*args, **kwargs):
            return func(*args, **kwargs)

        def wrapper(*args, **kwargs):
            if condition(*args, **kwargs):
                return cached_func(*args, **kwargs)
            else:
                return func(*args, **kwargs)

        return wrapper

    return decorator


if __name__ == "__main__":
    import argparse
    import uvicorn

    parser = argparse.ArgumentParser(
        description="Start the FastAPI server with command line arguments"
    )
    parser.add_argument(
        "--host", type=str, default="0.0.0.0", help="Host to run the server on"
    )
    parser.add_argument(
        "--port", type=int, default=8000, help="Port to run the server on"
    )
    parser.add_argument(
        "--reload", action="store_true", help="Enable auto-reload for development"
    )
    parser.add_argument("--compile", action="store_true", help="Enable model compile")
    parser.add_argument(
        "--lru_size",
        type=int,
        default=64,
        help="Set the size of the request cache pool, set it to 0 will disable lru_cache",
    )
    parser.add_argument(
        "--cors_origin",
        type=str,
        default="*",
        help="Allowed CORS origins. Use '*' to allow all origins.",
    )

    args = parser.parse_args()

    config.args = args

    if args.compile:
        print("Model compile is enabled")
        config.enable_model_compile = True

    def should_cache(*args, **kwargs):
        spk_seed = kwargs.get("spk_seed", -1)
        infer_seed = kwargs.get("infer_seed", -1)
        return spk_seed != -1 and infer_seed != -1

    if args.lru_size > 0:
        config.lru_size = args.lru_size
        generate.generate_audio = conditional_cache(should_cache)(
            generate.generate_audio
        )

    api = create_api()
    config.api = api

    if args.cors_origin:
        api.set_cors(allow_origins=[args.cors_origin])

    uvicorn.run(api.app, host=args.host, port=args.port, reload=args.reload)