"""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, AsrFilterResult
from asr_eventizer import AsrEventizer
from structure_builder import Structure, StructureBuilder
from structure_refiner import StructureRefiner
from markdown_builder import MarkdownBuilder
from windowizer import Windowizer

from agent import Agent
from utils import Timeline, 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"
    STRUCTURE_REFINER = "structure_refiner"
    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()
    ai_api_key = os.environ.get("AI_API_KEY", default=...)

    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(
        "--asr-ai-model",
        type=str,
        default="google/gemini-3.1-flash-lite"
    )
    parser.add_argument(
        "--asr-ai-base-url",
        type=str,
        default="https://api.proxyapi.ru/v1"
    )
    parser.add_argument(
        "--asr-ai-api-key",
        type=str,
        default=ai_api_key
    )
    parser.add_argument(
        "--structure-ai-model",
        type=str,
        default="google/gemini-3.5-flash"
    )
    parser.add_argument(
        "--structure-ai-base-url",
        type=str,
        default="https://api.proxyapi.ru/v1"
    )
    parser.add_argument(
        "--structure-ai-api-key",
        type=str,
        default=ai_api_key
    )
    parser.add_argument(
        "--refiner-ai-model",
        type=str,
        default="google/gemini-3.5-flash"
    )
    parser.add_argument(
        "--refiner-ai-base-url",
        type=str,
        default="https://api.proxyapi.ru/v1"
    )
    parser.add_argument(
        "--refiner-ai-api-key",
        type=str,
        default=ai_api_key
    )
    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't filter raw ASR output 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"))

def on_asr_events(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 eventizing")
        with open(WORKFLOW_DATA[current_step][0], "rb") as f:
            return (Step.VIDEO_REFERENCES, json.load(f))
    # bad request
    if input_data is None:
        logging.error("Can't create audio events without input_data")
        return (None, None)
    logging.info("Creating the agent")
    agent = Agent(
        model=ARGS.asr_ai_model,
        base_url=ARGS.asr_ai_base_url,
        api_key=ARGS.asr_ai_api_key
    )
    logging.info("Creating the windowizer")
    windowizer = Windowizer()
    logging.info("Creating audio events...")
    eventizer = AsrEventizer(agent, windowizer)
    result = eventizer.eventize(AsrFilterResult(**input_data))
    return (Step.VIDEO_REFERENCES, result.model_dump(mode="json"))

def on_video_references(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 video references")
        with open(WORKFLOW_DATA[current_step][0], "rb") as f:
            return (Step.REFERENCE_RESOLVER, json.load(f))
    if not os.path.isfile("video.mp4"):
        logging.info("No video, skipping")
        return (Step.REFERENCE_RESOLVER, None)
    # NOT IMPLEMENTED
    logging.warning("Video references are not implemented yet")
    return (Step.REFERENCE_RESOLVER, None)

def on_reference_resolver(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 reference resolver")
        with open(WORKFLOW_DATA[current_step][0], "rb") as f:
            return (Step.STRUCTURE_BUILDER, json.load(f))
    if not os.path.isfile("video.mp4"):
        logging.info("No video, skipping")
        return (Step.STRUCTURE_BUILDER, input_data)
    # NOT IMPLEMENTED
    logging.warning("Reference resolver is not implemented yet")
    return (Step.STRUCTURE_BUILDER, input_data)

def on_structure_builder(current_step: Step, input_data: dict | None) -> tuple[Step | None, dict | None]:
    # don't if done
    if os.path.isfile(WORKFLOW_DATA[current_step][0]):
        logging.info("Skipping structure builder")
        with open(WORKFLOW_DATA[current_step][0], "rb") as f:
            return (Step.STRUCTURE_REFINER, json.load(f))
    if input_data is None:
        logging.error("Can't build document structure without input_data")
        return (None, None)
    logging.info("Creating the agent")
    agent = Agent(
        model=ARGS.structure_ai_model,
        base_url=ARGS.structure_ai_base_url,
        api_key=ARGS.structure_ai_api_key
    )
    logging.info("Creating the windowizer")
    windowizer = Windowizer()
    logging.info("Building document structure...")
    builder = StructureBuilder(agent, windowizer)
    result = builder.build(Timeline(**input_data))
    return (Step.STRUCTURE_REFINER, result.model_dump(mode="json"))

def on_structure_refiner(current_step: Step, input_data: dict | None) -> tuple[Step | None, dict | None]:
    if os.path.isfile(WORKFLOW_DATA[current_step][0]):
        logging.info("Skipping structure refiner")
        with open(WORKFLOW_DATA[current_step][0], "rb") as f:
            return (Step.MARKDOWN_BUILDER, json.load(f))
    if input_data is None:
        logging.error("Can't refine document structure without input_data")
        return (None, None)
    logging.info("Creating the agent")
    agent = Agent(
        model=ARGS.refiner_ai_model,
        base_url=ARGS.refiner_ai_base_url,
        api_key=ARGS.refiner_ai_api_key
    )
    logging.info("Refining document structure...")
    refiner = StructureRefiner(agent, Windowizer())
    result = refiner.refine(Structure(**input_data))
    return (Step.MARKDOWN_BUILDER, result.model_dump(mode="json"))

def on_markdown_builder(current_step: Step, input_data: dict | None) -> tuple[Step | None, str | None]:
    if os.path.isfile(WORKFLOW_DATA[current_step][0]):
        logging.info("Skipping Markdown builder")
        return (None, None)
    if input_data is None:
        logging.error("Can't build Markdown without input_data")
        return (None, None)
    with open("events.json", "rb") as f:
        timeline = Timeline(**json.load(f))
    logging.info("Building Markdown...")
    builder = MarkdownBuilder(timeline)
    return (None, builder.build(Structure(**input_data)))

#
# Main
#
WORKFLOW_DATA: dict[Step, tuple[str, Callable[[Step, dict | None], tuple[Step | None, dict | str | 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", on_asr_events),
    Step.VIDEO_REFERENCES: ("unresolved.json", on_video_references),
    Step.REFERENCE_RESOLVER: ("events.json", on_reference_resolver),
    Step.STRUCTURE_BUILDER: ("structure.json", on_structure_builder),
    Step.STRUCTURE_REFINER: ("structure_refined.json", on_structure_refiner),
    Step.MARKDOWN_BUILDER: ("output.md", on_markdown_builder)
}
"""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:
            with open(output_file_path, "w", encoding="utf-8") as f:
                if isinstance(ret, str):
                    f.write(ret)
                else:
                    intermediate_result = ret
                    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()