import os
import sys
import time
import queue
import subprocess
import logging

import numpy as np
import cv2
import sounddevice as sd
from scipy.io import wavfile
from scipy.signal import resample
from PyQt6 import QtCore, QtGui, QtWidgets
from PyQt6.QtCore import pyqtSignal, QThread

import requests
import httpx
from bs4 import BeautifulSoup
from gne import GeneralNewsExtractor

# Import your custom models (ensure these are accessible in your PYTHONPATH)
from talkingface.audio_model import AudioModel
from talkingface.render_model import RenderModel

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class VideoWindow(QtWidgets.QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Video Window")
        self.setGeometry(100, 100, 800, 600)
        
        # Use layout for better scalability
        self.layout = QtWidgets.QVBoxLayout()
        self.setLayout(self.layout)
        
        self.label = QtWidgets.QLabel(self)
        self.label.setAlignment(QtCore.Qt.AlignmentFlag.AlignCenter)
        self.label.setScaledContents(True)  # Automatically scale image to fit QLabel size
        self.layout.addWidget(self.label)
        
        # Try to load a static image as the initial screen
        image_path = "subscribe.png"  # Ensure this image exists in the working directory
        test_image = cv2.imread(image_path)
        
        if test_image is not None:
            # Resize image to match QLabel dimensions (800x600)
            test_image = cv2.resize(test_image, (800, 600))
            logger.info(f"Loaded initial image from {image_path}.")
        else:
            # If loading fails, use a default green background
            logger.warning(f"Failed to load {image_path}. Using green screen instead.")
            test_image = np.zeros((600, 800, 3), dtype=np.uint8)
            test_image[:] = (0, 255, 0)  # Green background
        
        self.update_frame(test_image)

    @QtCore.pyqtSlot(np.ndarray)
    def update_frame(self, frame):
        if frame is None:
            logger.warning("Received None frame")
            return
        
        if frame.size == 0:
            logger.warning("Received empty frame")
            return
        
        try:
            # Ensure the incoming frame is in BGR format
            if len(frame.shape) == 3 and frame.shape[2] == 3:
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            else:
                frame_rgb = frame.copy()
            
            height, width, channel = frame_rgb.shape
            bytes_per_line = 3 * width
            
            # Create QImage object
            q_image = QtGui.QImage(
                frame_rgb.data,
                width,
                height,
                bytes_per_line,
                QtGui.QImage.Format.Format_RGB888
            )
            
            # Convert to QPixmap and set to QLabel
            pixmap = QtGui.QPixmap.fromImage(q_image)
            self.label.setPixmap(pixmap)
        except Exception as e:
            logger.error(f"Error updating frame: {e}")
            # Handle the exception appropriately

class VideoAudioWorker(QThread):
    frame_signal = pyqtSignal(np.ndarray)  # Signal emitted when a new frame is ready

    def __init__(self, parent=None):
        super().__init__(parent)
        self._run_flag = True

        # Initialize models
        self.audioModel = AudioModel()
        try:
            self.audioModel.loadModel("checkpoint/audio.pkl")
            logger.info("AudioModel loaded successfully.")
        except Exception as e:
            logger.error(f"Failed to load AudioModel: {e}")
            self._run_flag = False
            # Optionally, raise exception or handle it

        self.renderModel = RenderModel()
        try:
            self.renderModel.loadModel("checkpoint/render.pth")
            logger.info("RenderModel loaded successfully.")
        except Exception as e:
            logger.error(f"Failed to load RenderModel: {e}")
            self._run_flag = False
            # Optionally, raise exception or handle it

        # Set test video and keypoints
        self.test_video = "test3"
        self.pkl_path = f"video_data/{self.test_video}/keypoint_rotate.pkl"
        self.video_path = f"video_data/{self.test_video}/circle.mp4"
        try:
            self.renderModel.reset_charactor(self.video_path, self.pkl_path)
            logger.info("Character reset successfully.")
        except Exception as e:
            logger.error(f"Failed to reset character: {e}")
            self._run_flag = False
            # Optionally, raise exception or handle it

        # Audio processing
        self.wavpath = "output_audio.wav"
        self.wavpath1 = "output_audio1.wav"
        try:
            self.change_sample_rate(self.wavpath, self.wavpath1)
        except Exception as e:
            logger.error(f"Failed to change sample rate: {e}")
            self._run_flag = False
            # Optionally, raise exception or handle it

        try:
            self.rate, self.wav = wavfile.read(self.wavpath1, mmap=False)
            logger.info(f"Loaded WAV file with rate: {self.rate}")
        except Exception as e:
            logger.error(f"Failed to read WAV file: {e}")
            self.wav = None
            self._run_flag = False
            # Optionally, raise exception or handle it

        if self.wav is not None:
            # Normalize audio data
            if self.wav.dtype == np.int16:
                self.wav = self.wav.astype(np.float32) / 32768.0
            elif self.wav.dtype == np.int32:
                self.wav = self.wav.astype(np.float32) / 2147483648.0
            elif self.wav.dtype == np.uint8:
                self.wav = (self.wav.astype(np.float32) - 128) / 128.0
            elif self.wav.dtype == np.float32 or self.wav.dtype == np.float64:
                # Already in float format
                pass
            else:
                logger.error(f"Unsupported WAV format: {self.wav.dtype}")
                self._run_flag = False

            # Ensure sample rate is 16000
            self.sample_rate = 16000
            if self.rate != self.sample_rate:
                new_length = int(len(self.wav) * self.sample_rate / self.rate)
                if new_length > 0:
                    self.wav = resample(self.wav, new_length)
                    logger.info(f"Resampled WAV to {self.sample_rate} Hz.")
                else:
                    logger.error(f"Invalid resampling length: {new_length}")
                    self._run_flag = False
        else:
            logger.error("WAV data is not loaded properly.")
            self._run_flag = False

        if self._run_flag:
            # Set frame rate and samples per frame
            self.frame_rate = 25
            self.frame_duration = 1.0 / self.frame_rate
            self.samples_per_frame = int(self.sample_rate * self.frame_duration)
            self.index_ = 0

            # Initialize audio buffer
            self.audio_queue = queue.Queue(maxsize=20)  # Adjust size as needed

    def change_sample_rate(self, input_file, output_file, target_sample_rate=16000):
        if not os.path.isfile(input_file):
            raise FileNotFoundError(f"Input file {input_file} not found.")
        command = ['ffmpeg', '-y', '-i', input_file, '-ar', str(target_sample_rate), output_file]
        try:
            result = subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            logger.info(f"Sample rate changed to {target_sample_rate} Hz and saved as {output_file}.")
        except subprocess.CalledProcessError as e:
            error_msg = e.stderr.decode()
            logger.error(f"Error occurred during sample rate change: {error_msg}")
            raise

    def run(self):
        if not self._run_flag:
            logger.error("Initialization failed or models not loaded. Exiting thread.")
            return

        try:
            # Define audio callback function
            def audio_callback(outdata, frames, time_info, status):
                if status:
                    logger.warning(f"Stream status: {status}")
                try:
                    # Retrieve audio data from the queue
                    data = self.audio_queue.get_nowait()
                except queue.Empty:
                    # If the queue is empty, fill with silence
                    data = np.zeros(frames, dtype=np.float32)
                    logger.warning("Audio queue empty, filling with silence.")
                # Write data to the output buffer
                outdata[:, 0] = data

            kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])

            # Start the audio output stream
            with sd.OutputStream(samplerate=self.sample_rate, channels=1, dtype='float32',
                                 callback=audio_callback, blocksize=self.samples_per_frame):
                logger.info("Audio stream started.")

                # Record the start time
                start_time = time.perf_counter()

                while self._run_flag and self.index_ * self.samples_per_frame < len(self.wav):
                    # Calculate target time
                    target_time = start_time + self.index_ * self.frame_duration
                    current_time = time.perf_counter()
                    remaining_time = target_time - current_time

                    if remaining_time > 0:
                        time.sleep(remaining_time)
                    elif abs(remaining_time) > self.frame_duration:
                        logger.warning(f"Lagging by {-remaining_time:.2f} seconds at index {self.index_}")

                    frame_start_time = time.perf_counter()  # Record the start time of frame processing

                    start = self.index_ * self.samples_per_frame
                    end = start + self.samples_per_frame
                    if end > len(self.wav):
                        end = len(self.wav)
                    pcm_data = self.wav[start:end]

                    # Check if pcm_data is valid
                    if len(pcm_data) == 0:
                        logger.warning(f"No audio data at index {self.index_}.")
                        break  # Exit the loop if no data is left

                    # Process audio chunk
                    try:
                        mouth_frame = self.audioModel.interface_frame(pcm_data)
                        if mouth_frame is None:
                            raise ValueError("Received None from audioModel.interface_frame")
                        # Uncomment if filtering is needed
                        # mouth_frame = cv2.filter2D(mouth_frame, -1, kernel)
                        logger.debug(f"Processed audio chunk at index {self.index_}")
                    except Exception as e:
                        logger.error(f"Error processing audio at index {self.index_}: {e}")
                        mouth_frame = None

                    # Render video frame
                    try:
                        if mouth_frame is not None:
                            frame = self.renderModel.interface(mouth_frame)
                            if frame is not None:
                                # Uncomment if filtering is needed
                                # frame = cv2.bilateralFilter(frame, 9, 75, 75)
                                logger.debug(f"Generated frame at index {self.index_}")
                            else:
                                logger.warning(f"RenderModel returned None at index {self.index_}")
                        else:
                            frame = None
                    except Exception as e:
                        logger.error(f"Error generating frame at index {self.index_}: {e}")
                        frame = None

                    # Emit signal to main thread to update video frame
                    if frame is not None:
                        self.frame_signal.emit(frame)
                        logger.debug(f"Emitted frame at index {self.index_}")
                    else:
                        logger.warning(f"Skipping frame emission at index {self.index_} due to None frame.")

                    # Enqueue audio data
                    try:
                        self.audio_queue.put_nowait(pcm_data)
                        logger.debug(f"Enqueued audio chunk at index {self.index_}")
                    except queue.Full:
                        logger.warning("Audio queue is full, dropping audio chunk.")

                    # Update index
                    self.index_ += 1

                    # Calculate processing time and adjust sleep time
                    frame_end_time = time.perf_counter()
                    processing_time = frame_end_time - frame_start_time
                    sleep_time = self.frame_duration - processing_time
                    if sleep_time > 0:
                        time.sleep(sleep_time)
                    else:
                        logger.warning(f"Processing time {processing_time:.4f}s exceeds frame duration at index {self.index_ - 1}")

                logger.info("Audio stream finished.")
        except Exception as e:
            logger.error(f"Error in worker thread: {e}")
            # Optionally handle exception
        finally:
            logger.info("Worker thread exiting.")

    def stop(self):
        logger.info("Stopping worker thread...")
        self._run_flag = False
        self.wait()
        logger.info("Worker thread stopped.")

def fetch_html_content(url):
    """Fetch HTML content using httpx"""
    headers = {
        "User-Agent": "Mozilla/5.0"
    }
    
    try:
        response = httpx.get(url, headers=headers, timeout=10, proxies=None)
        if response.status_code == 200:
            return response.text
        else:
            return None
    except httpx.RequestError as e:
        logger.error(f"Request failed: {e}")
        return None

def extract_news_content(html):
    """Extract news content using GeneralNewsExtractor"""
    extractor = GeneralNewsExtractor()
    result = extractor.extract(html)
    return result

def getTTS(text):
    wave_path = "output_audio.wav"
    # Use getTTS1 to generate the audio
    getTTS1(text)

def getTTS1(text):
    url = "http://127.0.0.1:9880/tts"

    params = {
        "text": text,
        "text_lang": "zh",
        "ref_audio_path": "9.wav",
        "prompt_lang": "zh",
        "prompt_text": "我告诉你一个无比真实的事，差不多",
        "text_split_method": "cut5",
        "batch_size": 30,
        "media_type": "wav",
        "streaming_mode": "true"
    }
    try:
        response = requests.get(url, params=params, stream=True)
        response.raise_for_status()  # Raise an error for bad status codes

        # Save the audio file if the request is successful
        with open('output_audio.wav', 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        
        logger.info("Audio saved as output_audio.wav")

    except requests.exceptions.RequestException as e:
        logger.error(f"An error occurred: {e}")
        sys.exit(1)

def main():
 app = QtWidgets.QApplication(sys.argv)
 try:
    if len(sys.argv) < 2:
        print("Usage: python script.py <text_or_url>")
        sys.exit(1)

    href_value = sys.argv[1]
    logger.info(f"Received input: {href_value}")
   
    # Decide whether to treat input as text or URL
    if href_value.startswith('http://') or href_value.startswith('https://'):
        # Input is a URL, fetch and extract content
        html_content = fetch_html_content(href_value)
        if html_content:
            news_content = extract_news_content(html_content)
            text = news_content['content']
            logger.info(f"Extracted text from URL: {text}")
        else:
            logger.error("Failed to fetch or extract content from URL.")
            sys.exit(1)
    else:
        # Input is text
        text = href_value
        logger.info(f"Using input text: {text}")

    # Generate TTS audio
    #getTTS(text)

    # Create and display main window
    window = VideoWindow()
    window.show()

    # Create and start background thread
    worker = VideoAudioWorker()
    worker.frame_signal.connect(window.update_frame)
    worker.start()

    # Ensure the thread stops when the application exits
    app.aboutToQuit.connect(worker.stop)

    # Run the application
    sys.exit(app.exec())
 except Exception as e:
        print(f"Error in application: {str(e)}")
        
        os._exit(1)   # Exit th
 finally:
       os._exit(1) 

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        logger.error(f"Error in application: {str(e)}")
        sys.exit(1)