"""Application entry point"""

import traceback
import argparse
import logging
import json
import time
import os
from dataclasses import asdict
from enum import Enum
from typing import Callable

import torch
from asr import Asr, AsrRawResult
from asr_filter import AsrFilter

from agent import Agent
from utils import ffmpeg_split_video, ffmpeg_to_mp3

ARGS: argparse.Namespace

class Step(Enum):
    MEDIA_SEPARATION = "media_separation"
    VOICE_RECOGNITION = "voice_recognition"
    ASR_FILTER = "asr_filter"
    ASR_EVENTS = "asr_events"
    VIDEO_REFERENCES = "video_references"
    REFERENCE_RESOLVER = "reference_resolver"
    STRUCTURE_BUILDER = "structure_builder"
    MARKDOWN_BUILDER = "markdown_builder"

#
# Utility
#
def check_cuda() -> None:
    """Checks if CUDA is available."""
    if not torch.cuda.is_available():
        raise RuntimeError(
            "CUDA is unavaiable! Refusing to start, it's pointless."
        )

def setup_arguments() -> argparse.Namespace:
    """Parses CLI arguments and returns namespace. It may terminate the
    application on invalid arguments.
    
    Returns:
        - argparse namespace
    """
    voice_models = Asr.get_models_list()

    parser = argparse.ArgumentParser(
        prog="sumka",
        description="Summarizes large video/audio files into convenient format",
    )
    parser.add_argument(
        "--asr-model",
        choices=voice_models,
        default="turbo" if "turbo" in voice_models else voice_models[-1]
    )
    parser.add_argument(
        "--asr-language",
        choices=["ru", "en"],
        default="ru"
    )
    parser.add_argument(
        "--ai-model",
        type=str,
        default="GigaChat-3-Pro"
    )
    parser.add_argument(
        "--ai-base-url",
        type=str,
        default="https://api.giga.chat/v1"
    )
    parser.add_argument(
        "--ai-api-key",
        type=str
    )
    parser.add_argument("-v", action='store_true')
    return parser.parse_args()

#
# Workflow
#
def on_media_separation(current_step: Step, input_data: dict | None) -> tuple[Step | None, dict | None]:
    # there may be no data for this step
    if input_data:
        raise RuntimeError("There must be no input data for Media Separation")
    # do not split if there's `audio.mp3`
    if os.path.isfile(WORKFLOW_DATA[current_step][0]):
        logging.info("Skipping media separation")
        return (Step.VOICE_RECOGNITION, None)
    
    # filesnames to look for
    VIDEO_INPUTS = [
        "input.mp4",
        "input.mkv",
        "input.avi"
    ]
    AUDIO_INPUTS = [
        "input.mp3",
        "input.m4a",
        "input.wav"
    ]
    # execute video split
    for v in VIDEO_INPUTS:
        if os.path.isfile(v):
            logging.info(f"Splitting {v} to audio.mp3 and video.mp4")
            ffmpeg_split_video(v)
            return (Step.VOICE_RECOGNITION, None)
    # execute audio conversion
    for a in AUDIO_INPUTS:
        if os.path.isfile(a):
            logging.info(f"Converting {a} to audio.mp3")
            ffmpeg_to_mp3(a)
            return (Step.VOICE_RECOGNITION, None)
    logging.error(f"No supported `input.*` files found")
    return (None, None)

def on_voice_recognition(current_step: Step, input_data: dict | None) -> tuple[Step | None, dict | None]:
    # do not perform recognition if output file exists
    if os.path.isfile(WORKFLOW_DATA[current_step][0]):
        logging.info("Skipping voice recognition")
        with open(WORKFLOW_DATA[current_step][0], "rb") as f:
            return (Step.ASR_FILTER, json.load(f))
    logging.info(f"Loading ASR model `{ARGS.asr_model}`, language `{ARGS.asr_language}`")
    asr = Asr(ARGS.asr_model)
    logging.info(f"Speech recognition...")
    result = asr.recognize("audio.mp3", language=ARGS.asr_language)
    logging.info(f"Speech recognition done")
    return (Step.ASR_FILTER, result.model_dump(mode="json"))

def on_asr_filter(current_step: Step, input_data: dict | None) -> tuple[Step | None, dict | None]:
    # do not filter if output file exists
    if os.path.isfile(WORKFLOW_DATA[current_step][0]):
        logging.info("Skipping ASR filter")
        with open(WORKFLOW_DATA[current_step][0], "rb") as f:
            return (Step.ASR_EVENTS, json.load(f))
    # bad request
    if input_data is None:
        logging.error("Can'f filter raw ASR ouput without input_data")
        return (None, None)
    logging.info("Filtering raw ASR output...")
    filter = AsrFilter()
    result = filter.filter(AsrRawResult(**input_data))
    return (Step.ASR_EVENTS, result.model_dump(mode="json"))

#
# Main
#
WORKFLOW_DATA: dict[Step, tuple[str, Callable[[Step, dict | None], tuple[Step | None, dict | None]] | None]] = {
    Step.MEDIA_SEPARATION: ("audio.mp3", on_media_separation),
    Step.VOICE_RECOGNITION: ("asr_raw.json", on_voice_recognition),
    Step.ASR_FILTER: ("asr.json", on_asr_filter),
    Step.ASR_EVENTS: ("audio_events.json", None),
    Step.VIDEO_REFERENCES: ("unresolved.json", None),
    Step.REFERENCE_RESOLVER: ("events.json", None),
    Step.STRUCTURE_BUILDER: ("structure.json", None),
    Step.MARKDOWN_BUILDER: ("output.md", None)
}
"""Information about workflow.

Step.CODE: (
    "path/to/result.json",
    (cur_step: Step, input_data: dict | None)
        -> (next_step: Step | None, output_data: dict | None)
)
"""

def main() -> None:
    """Application entry point"""
    global ARGS
    check_cuda()
    ARGS = setup_arguments()
    logging.basicConfig(level=logging.DEBUG if ARGS.v else logging.INFO)
    step_to_do = Step.MEDIA_SEPARATION
    intermediate_result: dict | None = None
    # execute steps while possible
    while step_to_do:
        step_data = WORKFLOW_DATA[step_to_do]
        output_file_path = step_data[0]
        func = step_data[1]
        if not func:
            logging.error(f"Step {step_to_do} has no function, stopping")
            break
        logging.info(f"Executing step {step_to_do}")
        step_to_do, ret = func(step_to_do, intermediate_result)
        if ret is not None:
            intermediate_result = ret
            with open(output_file_path, "w") as f:
                json.dump(ret, f, ensure_ascii=False, indent=4)
            logging.info(f"Intermediate results are saved {output_file_path}")
    # final report
    logging.info(f"Workflow was interrupted at step {step_to_do}")

if __name__ == "__main__":
    try:
        main()
    except SystemExit:
        raise
    except:
        traceback.print_exc()