from math import pi
import os
import sys
from dataclasses import dataclass
import socket
from typing import List

now_dir = os.getcwd()
sys.path.append(now_dir)
from scipy.io import wavfile

from configs.config import Config
from infer.modules.vc.modules import VC

import torch
import json
import traceback
import numpy as np


@dataclass
class Invocation:
    f0up_key: int
    input_path: str
    opt_path: str
    model_name: str
    index_rate: float
    filter_radius: int
    rms_mix_rate: float
    protect: float
    pitch_curve: List[List[float]]
    rundir: str

    @classmethod
    def from_json(cls, j):
        return cls(
            j["f0up_key"],
            j["input_path"],
            j["opt_path"],
            j["model_name"],
            j["index_rate"],
            j["filter_radius"],
            j["rms_mix_rate"],
            j["protect"],
            j["pitch_curve"],
            j["rundir"],
        )


def recvall(sock):
    data = bytearray()
    while True:
        packet = sock.recv(4096)
        if not packet:
            return data
        data.extend(packet)


current_model_name = ""
current_index_path = ""


def handle_conn(vc, conn):
    buffer = recvall(conn)
    args = Invocation.from_json(json.loads(buffer.decode("utf-8")))

    global current_model_name, current_index_path
    if current_model_name != args.model_name:
        vc.get_vc("")  # clear cache
        loaded = vc.get_vc(args.model_name, relative_to=args.rundir)[0]

        current_model_name = args.model_name
        current_index_path = loaded["index_path"]

    _, wav_opt = vc.vc_single(
        0,
        args.input_path,
        args.f0up_key,
        "rmvpe",
        current_index_path,
        None,
        args.index_rate,
        args.filter_radius,
        0,
        args.rms_mix_rate,
        args.protect,
        np.array(args.pitch_curve) if args.pitch_curve is not None else None,
    )
    wavfile.write(args.opt_path, wav_opt[0], wav_opt[1])

    conn.sendall(
        json.dumps({"status": "ok", "pitch_curve": wav_opt[2].tolist()}).encode("utf-8")
    )


def main():
    config = Config()
    config.device = "cuda" if torch.cuda.is_available() else "cpu"
    vc = VC(config)
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("localhost", 12345))
        s.listen()
        while True:
            conn = None
            try:
                conn, addr = s.accept()
                handle_conn(vc, conn)
            except Exception as e:
                if conn is None:
                    raise e
                try:
                    traceback.print_exception(e)
                    conn.sendall(
                        json.dumps({"status": "error", "message": str(e)}).encode(
                            "utf-8"
                        )
                    )
                    conn.shutdown(socket.SHUT_RDWR)
                except Exception as e2:
                    print("Failed to send exception to client: " + str(e2))
            finally:
                if conn is not None:
                    conn.close()


if __name__ == "__main__":
    main()