from concurrent.futures import ThreadPoolExecutor
import glob
import os
from pathlib import Path
import shutil
import signal
import socket
import sys
import tempfile
import threading
import time
import trace
import traceback
import boto3
from numpy import Inf
from lib.redis_client import RedisClient, getenv
import uuid
import subprocess
import json
import zipfile
from lib.progress import ProgressHandler
from lib.audio_utils import download_audio, upload_audio
import psycopg
from psycopg.rows import dict_row
import backoff


SAMPLE_RATE = 48000  # rate used for rvc

MAX_CACHE_SIZE = 4 * 1024 * 1024 * 1024  # 4GB


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


RVC_MODELS_S3_BUCKET = getenv("RVC_MODELS_S3_BUCKET")
RVC_SAMPLES_S3_BUCKET = getenv("RVC_SAMPLES_S3_BUCKET")
RVC_CACHE_DIR = getenv("RVC_CACHE_DIR")

os.makedirs(RVC_CACHE_DIR, exist_ok=True)


def run_with_line_callback_thd(popen_handle, callback):
    pipe = popen_handle.stdout

    def terminate():
        try:
            os.killpg(popen_handle.pid, signal.SIGINT)
        except Exception as e:
            print(f"Failed to terminate subprocess process group: {e}")

    while True:
        line = pipe.readline()
        if not line:
            break
        try:
            sys.stdout.write(line)
            callback(line, terminate)
        except Exception as e:
            print(f"line callback failed: {e}")


def run_with_line_callback(args, callback):
    p = subprocess.Popen(
        args,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )
    thd = threading.Thread(
        target=run_with_line_callback_thd,
        args=(p, callback),
        daemon=True,
    )
    try:
        thd.start()
        p.wait()
        if p.returncode != 0:
            raise RuntimeError(f"Subprocess exited with code {p.returncode}")
    finally:
        if p.poll() is None:
            p.kill()
        p.wait()
        try:
            os.killpg(p.pid, signal.SIGINT)
        except Exception as e:
            print(f"Failed to terminate subprocess process group: {e}")
        if thd.is_alive():
            thd.join()


class RVCTrainer:
    def __init__(self, models_s3_bucket, samples_s3_bucket, cache_dir):
        self.s3 = boto3.client("s3", verify=False)
        self.models_s3_bucket = models_s3_bucket
        self.samples_s3_bucket = samples_s3_bucket
        self.cache_dir = cache_dir

    @classmethod
    def from_env(cls):
        return cls(RVC_MODELS_S3_BUCKET, RVC_SAMPLES_S3_BUCKET, RVC_CACHE_DIR)

    def _execute(self, query, args, fetchall=False):
        # pylint: disable=E1129
        with psycopg.connect("", row_factory=dict_row, connect_timeout=10) as conn:
            conn.autocommit = True
            with conn.cursor() as cur:
                cur.execute(query, args)
                if fetchall:
                    return cur.fetchall()

    @backoff.on_exception(
        backoff.fibo,
        (
            psycopg.OperationalError,
            psycopg.InterfaceError,
            psycopg.errors.ConnectionTimeout,
            OSError,
        ),
        max_time=60,
    )
    def _execute_backoff(self, query, args, fetchall=False):
        return self._execute(query, args, fetchall=fetchall)

    def run(self):
        # pylint: disable=E1129
        with psycopg.connect("", row_factory=dict_row, connect_timeout=10) as conn:
            conn.autocommit = True
            with conn.transaction(), conn.cursor() as cur:
                cur.execute(
                    """
                    SELECT r.id, vcm.uuid model_uuid, COALESCE(corpus_samples.sample_s3_keys, '{}') corpus_s3_keys
                    FROM voice_conversion_model_training_runs r
                    JOIN voice_conversion_models vcm ON vcm.id = r.model_id
                    LEFT JOIN LATERAL (
                        SELECT ARRAY_AGG(DISTINCT COALESCE(s.s3_key, s.compressed_s3_key)) sample_s3_keys FROM voice_conversion_model_corpus c
                        JOIN samples s ON s.id = c.sample_id
                        WHERE c.model_id = vcm.id
                    ) corpus_samples ON true
                    WHERE NOT vcm.archived AND r.state = 'pending'
                    ORDER BY r.created_at ASC LIMIT 1
                    FOR UPDATE OF r SKIP LOCKED
                    """
                )
                rows = cur.fetchall()
                if len(rows) == 0:
                    print("No pending training runs - exiting")
                    sys.exit(0)
                row = rows[0]
                training_run_id = row["id"]
                model_uuid = row["model_uuid"]
                corpus_s3_keys = row["corpus_s3_keys"]
                cur.execute(
                    "UPDATE voice_conversion_model_training_runs SET (state, trainer_version) = ('training', 'v0') WHERE id = %s",
                    (training_run_id,),
                )

            try:
                with open("/trainer_working", "w") as f:
                    f.write("ready")
            except (PermissionError, FileNotFoundError):
                print(f"Failed to write /trainer_working")

            def notify_progress(progress):
                try:
                    self._execute(
                        "UPDATE voice_conversion_model_training_runs SET progress = %s WHERE id = %s",
                        (progress, training_run_id),
                    )
                except Exception as e:
                    print(f"Failed to update progress: {e}")

            def poll_cancellation():
                try:
                    rows = self._execute(
                        "SELECT 1 FROM voice_conversion_model_training_runs WHERE id = %s AND state = 'training'",
                        (training_run_id,),
                        fetchall=True,
                    )
                    return len(rows) == 0
                except Exception as e:
                    print(f"Failed to poll cancellation: {e}")
                    return False

            try:
                s3_key = self.handle_training_request(
                    corpus_s3_keys, model_uuid, notify_progress, poll_cancellation
                )
            except Exception as e:
                formatted_exn = traceback.format_exception(e)
                print(formatted_exn)
                self._execute_backoff(
                    "UPDATE voice_conversion_model_training_runs SET (state, error) = ('failed', %s) WHERE id = %s",
                    (formatted_exn, training_run_id),
                )
                return

            self._execute_backoff(
                "UPDATE voice_conversion_model_training_runs SET (state, s3_key) = ('finished', %s) WHERE id = %s",
                (s3_key, training_run_id),
            )

    def handle_training_request(
        self, corpus_s3_keys, model_uuid, notify_progress, poll_cancellation
    ):
        with tempfile.TemporaryDirectory() as tmpdir, ThreadPoolExecutor(
            max_workers=4
        ) as executor:
            notify_progress(1.0)
            download_futures = [
                executor.submit(
                    download_audio,
                    self.s3,
                    self.samples_s3_bucket,
                    s3_key,
                    tmpdir,
                    sample_rate=SAMPLE_RATE,
                )
                for s3_key in corpus_s3_keys
            ]
            for future in download_futures:
                future.result()
            notify_progress(2.0)
            epoch_count = 0
            epoch_total = 100.0

            def progress_callback(line, terminate):
                nonlocal epoch_total
                nonlocal epoch_count
                if line.startswith("INFO:rvc_model:====> Epoch:"):
                    epoch_count += 1
                    if poll_cancellation():
                        terminate()
                    notify_progress(2.0 + epoch_count / epoch_total * 97.0)
                elif line.startswith("************ Training for "):
                    try:
                        epoch_total = float(line.split(" ")[3])
                    except Exception:
                        print("Failed to parse total epoch count")

            with tempfile.TemporaryDirectory() as rundir:
                run_with_line_callback(
                    [
                        "python",
                        "train-cli.py",
                        tmpdir,
                        "rvc_model",
                        rundir,
                    ],
                    progress_callback,
                )
                model_uuid = uuid.uuid4()
                s3_key = f"{model_uuid}.zip"
                with tempfile.TemporaryDirectory(
                    dir=self.cache_dir, ignore_cleanup_errors=True
                ) as pack_dir:
                    shutil.move(
                        f"{rundir}/assets/weights/rvc_model.pth",
                        f"{pack_dir}/rvc_model.pth",
                    )
                    for path in glob.glob(f"{rundir}/logs/rvc_model/*.index"):
                        shutil.move(path, f"{pack_dir}/{os.path.basename(path)}")
                    with zipfile.ZipFile(f"{tmpdir}/{s3_key}", "w") as zip:
                        for path in glob.glob(f"{pack_dir}/*"):
                            zip.write(path, os.path.basename(path))
                    shutil.move(pack_dir, f"{self.cache_dir}/{model_uuid}")
                self.s3.upload_file(f"{tmpdir}/{s3_key}", self.models_s3_bucket, s3_key)
                notify_progress(100.0)
                return s3_key


# hack to make local training work
class RVCRedisClientLocalTraining(RedisClient):
    def __init__(self, models_s3_bucket, samples_s3_bucket, cache_dir, *args):
        super().__init__(*args)
        self.s3 = boto3.client("s3", verify=False)
        self.models_s3_bucket = models_s3_bucket
        self.samples_s3_bucket = samples_s3_bucket
        self.cache_dir = cache_dir

    @classmethod
    def from_env(cls):
        return super().from_env(
            RVC_MODELS_S3_BUCKET,
            RVC_SAMPLES_S3_BUCKET,
            RVC_CACHE_DIR,
        )

    def train(self, corpus_s3_keys, model_uuid):
        r = RVCTrainer.from_env()

        def notify_progress(progress):
            self.client.rpush(
                f"rvc_progress:{model_uuid}",
                json.dumps(
                    {
                        "state": "training",
                        "progress": progress,
                    }
                ),
            )

        def poll_cancellation():
            return self.client.lpop(f"rvc_cancel:{model_uuid}") is not None

        try:
            s3_key = r.handle_training_request(
                corpus_s3_keys, model_uuid, notify_progress, poll_cancellation
            )
            self.client.rpush(
                f"rvc_progress:{model_uuid}",
                json.dumps(
                    {
                        "state": "finished",
                        "s3Key": s3_key,
                    }
                ),
            )
        except Exception as e:
            formatted_exn = traceback.format_exception(e)
            print(formatted_exn)
            self.client.rpush(
                f"rvc_progress:{model_uuid}",
                json.dumps(
                    {
                        "state": "failed",
                        "error": formatted_exn,
                    }
                ),
            )


def dirsize(path):
    total_size = 0
    for dirpath, _, filenames in os.walk(path):
        for f in filenames:
            fp = os.path.join(dirpath, f)
            total_size += os.path.getsize(fp)
    return total_size


class RVCRedisClient(RedisClient):
    def __init__(self, models_s3_bucket, samples_s3_bucket, cache_dir, *args):
        super().__init__(*args)
        self.s3 = boto3.client("s3", verify=False)
        self.models_s3_bucket = models_s3_bucket
        self.samples_s3_bucket = samples_s3_bucket
        self.cache_dir = cache_dir
        self.inference_proc = subprocess.Popen(
            ["python", "tools/infer_socket.py"],
        )
        self.inference_watchdog = threading.Thread(
            target=self._die_when_subprocess_exits,
            args=(self.inference_proc,),
            daemon=True,
        )
        self.inference_watchdog.start()

    @classmethod
    def from_env(cls):
        return super().from_env(
            RVC_MODELS_S3_BUCKET,
            RVC_SAMPLES_S3_BUCKET,
            RVC_CACHE_DIR,
        )

    def _die_when_subprocess_exits(self, proc):
        proc.wait()
        os.kill(os.getpid(), signal.SIGINT)

    def _cleanup_cache(self):
        cache_size = dirsize(self.cache_dir)
        for _ in range(5):  # don't work too hard
            if cache_size / MAX_CACHE_SIZE < 0.8:
                return
            items = sorted(
                (os.path.getmtime(f"{self.cache_dir}/{f}"), f)
                for f in os.listdir(self.cache_dir)
            )
            for _, f in items:
                print(f"cache: removing {f}")
                if os.path.isdir(f"{self.cache_dir}/{f}"):
                    cache_size -= dirsize(f"{self.cache_dir}/{f}")
                    shutil.rmtree(f"{self.cache_dir}/{f}")
                else:
                    cache_size -= os.path.getsize(f"{self.cache_dir}/{f}")
                    os.unlink(f"{self.cache_dir}/{f}")

    def _download_model(self, s3_key, notify_progress):
        progress_handler = ProgressHandler(notify_progress, phase="Downloading model")

        try:
            self._cleanup_cache()
        except Exception as e:
            print(f"Failed to cleanup cache: {e}")

        model_uuid = s3_key.split(".")[0]
        dest_dir = f"{self.cache_dir}/{model_uuid}"
        if os.path.exists(dest_dir):
            progress_handler(1.0)
            os.utime(dest_dir)
            return dest_dir
        with tempfile.TemporaryDirectory() as tmpdir, tempfile.TemporaryDirectory(
            dir=self.cache_dir, ignore_cleanup_errors=True
        ) as unpack_dir:
            meta_data = self.s3.head_object(Bucket=self.models_s3_bucket, Key=s3_key)
            total_size = meta_data.get("ContentLength", None)
            bytes_downloaded = 0
            last_notify_time = 0

            def s3_progress_callback(chunk):
                nonlocal bytes_downloaded, last_notify_time
                bytes_downloaded += chunk
                if total_size is None:
                    return
                if time.time() - last_notify_time < 0.5:
                    return
                last_notify_time = time.time()
                progress_handler(bytes_downloaded / total_size)

            self.s3.download_file(
                self.models_s3_bucket,
                s3_key,
                f"{tmpdir}/{s3_key}",
                Callback=s3_progress_callback,
            )
            with zipfile.ZipFile(f"{tmpdir}/{s3_key}", "r") as z:
                z.extractall(unpack_dir)
            shutil.move(unpack_dir, dest_dir)
        progress_handler(1.0)
        return dest_dir

    def warmup(self):
        for i in range(2):
            with tempfile.TemporaryDirectory() as rundir:
                os.makedirs(f"{rundir}/assets/weights", exist_ok=True)
                shutil.copy(
                    f"./assets/weights/warmup.pth",
                    f"{rundir}/assets/weights/warmup.pth",
                )
                os.makedirs(f"{rundir}/logs/warmup", exist_ok=True)
                for path in glob.glob(f"./logs/warmup/*.index"):
                    shutil.copy(path, f"{rundir}/logs/warmup/{os.path.basename(path)}")
                self._run_inference(
                    0,
                    "assets/warmup.wav",
                    "/tmp/deleteme.wav",
                    "warmup.pth",
                    None if i == 0 else [[0.0, 400.0], [4.0, 500.0]],
                    rundir,
                )
        print("warmup ok")

    def _run_inference(
        self,
        f0up_key,
        input_path,
        output_path,
        model_name,
        input_pitch_curve,
        rundir,
    ):
        request = {
            "f0up_key": f0up_key,
            "input_path": input_path,
            "opt_path": output_path,
            "model_name": model_name,
            "index_rate": 0.66,
            "filter_radius": 3,
            "rms_mix_rate": 1,
            "protect": 0.33,
            "pitch_curve": input_pitch_curve,
            "rundir": rundir,
        }
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            while True:
                try:
                    s.connect(("localhost", 12345))
                    break
                except ConnectionRefusedError:
                    time.sleep(1)
            s.sendall(json.dumps(request).encode("utf-8"))
            s.shutdown(socket.SHUT_WR)
            buffer = recvall(s)
            response = json.loads(buffer.decode("utf-8"))
            if "status" not in response:
                raise RuntimeError("No status in response")
            if response["status"] == "error":
                raise RuntimeError(response["message"])

        if not os.path.exists(output_path):
            raise RuntimeError("Output file not found")

        return response["pitch_curve"]

    def check_pitch_curve(self, pitch_curve):
        assert pitch_curve is None or (
            isinstance(pitch_curve, list)
            and all(
                isinstance(item, list)
                and len(item) == 2
                and all(isinstance(x, float) for x in item)
                for item in pitch_curve
            )
            and all(
                pitch_curve[i][0] < pitch_curve[i + 1][0]
                for i in range(len(pitch_curve) - 1)
            )
        ), "pitch curve must be a list of pairs of floats (t,p) where t is increasing"
        return pitch_curve

    def handle_request(self, request_data, recorder, notify_progress, times_out_at):
        sample_s3_key = request_data["sampleS3Key"]
        model_s3_key = request_data["modelS3Key"]
        duration = max(0.1, request_data.get("sampleDuration", 10.0))
        codec = "ogg"
        input_pitch_curve = self.check_pitch_curve(request_data.get("pitchCurve", None))
        model_uuid = model_s3_key.split(".")[0]
        f0up_key = int(request_data.get("f0Delta", 0))
        model_dir = self._download_model(model_s3_key, notify_progress)
        with tempfile.TemporaryDirectory() as rundir:
            os.makedirs(f"{rundir}/assets/weights", exist_ok=True)
            shutil.copy(
                f"{model_dir}/rvc_model.pth",
                f"{rundir}/assets/weights/{model_uuid}.pth",
            )
            os.makedirs(f"{rundir}/logs/{model_uuid}", exist_ok=True)
            for path in glob.glob(f"{model_dir}/*.index"):
                shutil.copy(
                    path, f"{rundir}/logs/{model_uuid}/{os.path.basename(path)}"
                )

            input_path = download_audio(
                self.s3,
                self.samples_s3_bucket,
                sample_s3_key,
                rundir,
                progress_handler=ProgressHandler(
                    notify_progress, duration=duration, phase="Decoding"
                ),
                sample_rate=SAMPLE_RATE,
            )
            output_path = f"{rundir}/out.wav"
            inference_start_time = time.time()
            output_pitch_curve = self._run_inference(
                f0up_key,
                input_path,
                output_path,
                f"{model_uuid}.pth",
                input_pitch_curve,
                rundir,
            )
            inference_duration = time.time() - inference_start_time
            print(f"Inference took {inference_duration} seconds")

            output_s3_key = f"{uuid.uuid4()}.{codec}"
            upload_audio(
                self.s3,
                self.samples_s3_bucket,
                output_s3_key,
                output_path,
                codec,
                progress_handler=ProgressHandler(
                    notify_progress, duration=duration, phase="Encoding"
                ),
            )
            return {
                "state": "finished",
                "s3Key": output_s3_key,
                "pitchCurve": output_pitch_curve,
            }


if __name__ == "__main__":
    if len(sys.argv) > 1:
        if sys.argv[1] == "--train":
            client = RVCTrainer.from_env()
        elif sys.argv[1] == "--local-train":
            trainer = RVCRedisClientLocalTraining.from_env()
            trainer.train(sys.argv[2].split(","), sys.argv[3])
            sys.exit(0)
        else:
            raise RuntimeError("Unknown argument")
    else:
        client = RVCRedisClient.from_env()
    client.run()