import cv2
import easyocr
import pysrt
import numpy as np
from tqdm import tqdm
from dataclasses import dataclass

# =============================
# Language Priority Mapping
# =============================
OCR_LANG_MAP = {
    "english": ["en"],
    "japanese": ["ja"],
    "korean": ["ko"],
    "chinese_simplified": ["ch_sim", "en"],
    "chinese_traditional": ["ch_tra", "en"],
    "thai": ["th"]
}

# =============================
# Data Models
# =============================
@dataclass
class SubtitleBox:
    x: int
    y: int
    width: int
    height: int

@dataclass
class SubtitleLine:
    text: str
    start: float
    end: float
    status: str  # entered | active | exited

# =============================
# Video Reader (Frame Accurate)
# =============================
class VideoReader:
    def __init__(self, path):
        self.cap = cv2.VideoCapture(path)
        if not self.cap.isOpened():
            raise Exception("Cannot open video file")

        self.fps = self.cap.get(cv2.CAP_PROP_FPS)
        self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))

    def time_to_frame(self, t):
        return int(t * self.fps)

    def get_frame(self, index):
        self.cap.set(cv2.CAP_PROP_POS_FRAMES, index)
        ret, frame = self.cap.read()
        return frame if ret else None

# =============================
# OCR Engine
# =============================
class SubtitleOCR:
    def __init__(self, language_key):
        langs = OCR_LANG_MAP[language_key]
        self.reader = easyocr.Reader(langs, gpu=False)

    def recognize(self, image):
        result = self.reader.readtext(image, detail=0)
        return " ".join(result).strip()

# =============================
# Subtitle Tracker
# =============================
class SubtitleTracker:
    def __init__(self):
        self.last_text = ""
        self.start_time = None
        self.subtitles = []

    def update(self, text, current_time):
        if text != self.last_text:
            if self.last_text.strip() != "":
                self.subtitles.append(
                    SubtitleLine(
                        text=self.last_text,
                        start=self.start_time,
                        end=current_time,
                        status="exited"
                    )
                )
            self.last_text = text
            self.start_time = current_time
            return "entered"
        return "active"

    def finalize(self, end_time):
        if self.last_text.strip() != "":
            self.subtitles.append(
                SubtitleLine(
                    text=self.last_text,
                    start=self.start_time,
                    end=end_time,
                    status="exited"
                )
            )

# =============================
# Crop Subtitle Region
# =============================
def crop_subtitle(frame, box: SubtitleBox):
    return frame[
        box.y : box.y + box.height,
        box.x : box.x + box.width
    ]

# =============================
# Main OCR Process
# =============================
def process_video(
    video_path,
    subtitle_box,
    start_time,
    end_time,
    language_key
):
    video = VideoReader(video_path)
    ocr = SubtitleOCR(language_key)
    tracker = SubtitleTracker()

    start_frame = video.time_to_frame(start_time)
    end_frame = video.time_to_frame(end_time)

    for frame_idx in tqdm(range(start_frame, end_frame)):
        frame = video.get_frame(frame_idx)
        if frame is None:
            continue

        crop = crop_subtitle(frame, subtitle_box)
        text = ocr.recognize(crop)

        current_time = frame_idx / video.fps
        tracker.update(text, current_time)

    tracker.finalize(end_time)
    return tracker.subtitles

# =============================
# Export SRT
# =============================
def export_srt(subtitles, output_path):
    srt = pysrt.SubRipFile()

    for i, sub in enumerate(subtitles, start=1):
        srt.append(
            pysrt.SubRipItem(
                index=i,
                start=pysrt.SubRipTime(seconds=sub.start),
                end=pysrt.SubRipTime(seconds=sub.end),
                text=sub.text
            )
        )

    srt.save(output_path, encoding="utf-8")

# =============================
# Example Run
# =============================
if __name__ == "__main__":
    video_path = "input.mp4"

    subtitle_box = SubtitleBox(
        x=200,     # Horizontal Position
        y=800,     # Vertical Position
        width=1500,
        height=200
    )

    start_time = 10.0
    end_time = 120.0

    language = "chinese_simplified"

    subtitles = process_video(
        video_path,
        subtitle_box,
        start_time,
        end_time,
        language
    )

    export_srt(subtitles, "output.srt")
    print("OCR Finished → output.srt")