import os
import sys
import subprocess
import threading
import time
import requests
import socket
import io
import wave
import json
import re
import numpy as np
import yt_dlp
from flask import Flask, render_template, request
from flask_socketio import SocketIO, emit
from dotenv import load_dotenv

from google import genai
from google.genai import types
from groq import Groq

# DNSパッチ
_original_getaddrinfo = socket.getaddrinfo

def patched_getaddrinfo(host, port, family=0, type=0, proto=0, flags=0):
    if host and isinstance(host, str) and ("youtube" in host or "youtu.be" in host or "googlevideo" in host):
        try:
            return _original_getaddrinfo(host, port, family, type, proto, flags)
        except Exception:
            try:
                url = f"https://8.8.8.8/resolve?name={host}&type=A"
                r = requests.get(url, timeout=2)
                data = r.json()
                if "Answer" in data:
                    ip = data["Answer"][0]["data"]
                    return _original_getaddrinfo(ip, port, family, type, proto, flags)
            except Exception:
                pass
    return _original_getaddrinfo(host, port, family, type, proto, flags)

socket.getaddrinfo = patched_getaddrinfo

load_dotenv()

# Cookie読み込み
try:
    cookie_content = os.getenv("YOUTUBE_COOKIES")
    if cookie_content:
        if "\\n" in cookie_content:
            cookie_content = cookie_content.replace("\\n", "\n")
        if not cookie_content.endswith("\n"):
            cookie_content += "\n"
        with open("cookies.txt", "w") as f:
            f.write(cookie_content)
except Exception:
    pass

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
model_name = "gemini-2.0-flash"

if not GEMINI_API_KEY:
    print("⚠️ GEMINI_API_KEYが設定されていません")

groq_client = None
if GROQ_API_KEY:
    try:
        groq_client = Groq(api_key=GROQ_API_KEY)
        print("✅ Groq Client Initialized")
    except Exception:
        pass

sys.stdout.reconfigure(line_buffering=True)

app = Flask(__name__)
app.config['SECRET_KEY'] = os.getenv("SECRET_KEY", "secret!")
socketio = SocketIO(app, cors_allowed_origins="*", async_mode='threading')

MAX_CONCURRENT_STREAMS = 10
SPEAKER_COLORS = ["#FFFFFF", "#03DAC6", "#FFB400", "#CF6679", "#BB86FC"]

# Whisperが無音時によく出力してしまう無意味なフレーズ集
HALLUCINATION_PHRASES = [
    "You", "Thank you", "Thank you.", "Thanks for watching", "Bye",
    "Okay", "Okay.", "All right", "music", "[Music]", "Subscribe",
    "Copyright", "MBC", "SBS", "news", "Amara", "Subtitles by",
    "Caption", "Captions", "Provided by", "http", "www", ".com",
    "Silence", "[Silence]", "Audio", "Sound", "noise"
]

session_speakers = {}
active_sessions = {}

# ログ抑制用
class QuietLogger:
    def debug(self, msg): pass
    def warning(self, msg): pass
    def error(self, msg): pass

class StreamWorker:
    def __init__(self, sid, url, settings, socket_instance):
        self.sid = sid
        self.url = url
        self.settings = settings
        self.socketio = socket_instance
        self.is_running = False
        self.ffmpeg_proc = None
        self.fifo_path = f"/tmp/stream_{sid}.fifo"
        
        # フィラー（言い淀み）の削除と、要約的な翻訳を指示
        self.system_instruction = """
        You are a professional simultaneous interpreter.
        
        INPUT DATA:
        1. Transcript: English text transcribed from audio.
        2. Audio: Source sound clip.
        
        YOUR TASK:
        1. Identify the speaker (Speaker 1, Speaker 2...).
        2. Translate the Transcript into natural, concise Japanese.
        
        CRITICAL RULES for Translation:
        - **REMOVE FILLERS**: Ignore words like "uh", "um", "like", "you know", "I mean", "actually".
        - **IGNORE REPETITIONS**: If the user stutters (e.g., "I- I- I went"), translate only the final meaning ("I went").
        - **NO HALLUCINATIONS**: If the input text is just "Thank you" or "[Music]" or random noise, output NOTHING (empty string).
        - **CONCISE**: Do not translate word-for-word. Capture the *intent*.
        - DO NOT include speaker labels.
        - Output ONLY the translated text.
        """
        self.client = genai.Client(api_key=GEMINI_API_KEY)
        self.model_name = model_name
        # --- 擬似話者識別用 ---
        self.current_speaker = "Speaker 1"
        self.last_rms = 0.0

    def start(self):
        if self.is_running: return
        self.is_running = True
        self.socketio.start_background_task(target=self._main_loop)

    def stop(self):
        self.is_running = False
        if self.ffmpeg_proc and self.ffmpeg_proc.poll() is None:
            self.ffmpeg_proc.kill()
        if os.path.exists(self.fifo_path):
            try: os.remove(self.fifo_path)
            except: pass
        print(f"🛑 Worker停止: {self.sid}")

    def _run_ytdlp(self):
        if os.path.exists(self.fifo_path):
            os.remove(self.fifo_path)
        os.mkfifo(self.fifo_path)

        base_opts = {
            'format': 'bestaudio/best',
            'outtmpl': self.fifo_path,
            'quiet': True,
            'no_warnings': True,
            'logger': QuietLogger(),
            'nocheckcertificate': True,
            'socket_timeout': 30,
            'retries': 10,
            'fragment_retries': 10,
            'force_ipv4': False,
            'live_from_start': False,
            'wait_for_video': (1, 15),
            'ignoreerrors': True,
            'cachedir': False,
        }

        strategies = []
        if os.path.exists("cookies.txt"):
            opts_cookie = base_opts.copy()
            opts_cookie['cookiefile'] = 'cookies.txt'
            opts_cookie['extractor_args'] = {'youtube': {'impersonate': ['chrome']}}
            strategies.append(("Cookie(Web)", opts_cookie))
        
        opts_ios = base_opts.copy()
        opts_ios['cookiefile'] = None
        opts_ios['extractor_args'] = {'youtube': {'player_client': ['ios']}}
        strategies.append(("NoCookie(iOS)", opts_ios))
        
        opts_tv = base_opts.copy()
        opts_tv['cookiefile'] = None
        opts_tv['extractor_args'] = {'youtube': {'player_client': ['tv_embedded']}}
        strategies.append(("NoCookie(TV)", opts_tv))

        success = False
        for label, opts in strategies:
            if not self.is_running: break
            try:
                with yt_dlp.YoutubeDL(opts) as ydl:
                    ydl.download([self.url])
                if os.path.exists(self.fifo_path):
                    success = True
                    break
            except Exception:
                time.sleep(1)
                continue

        if not success:
            if self.is_running: 
                self.socketio.emit("error_msg", {"msg": "動画にアクセスできませんでした。"}, room=self.sid)
            self.is_running = False

    def _create_wav_bytes(self, pcm_data):
        wav_io = io.BytesIO()
        try:
            with wave.open(wav_io, 'wb') as wb:
                wb.setnchannels(1)
                wb.setsampwidth(2)
                wb.setframerate(16000)
                wb.writeframes(pcm_data)
            wav_io.seek(0)
            return wav_io
        except Exception:
            return None

    def _transcribe_groq(self, audio_data):
        if not groq_client: return None
        try:
            audio_file = self._create_wav_bytes(audio_data)
            if not audio_file: return None
            
            audio_file.name = "audio.wav"
            transcription = groq_client.audio.transcriptions.create(
                file=audio_file,
                model="whisper-large-v3",
                response_format="text",
                language="en",
                temperature=0.0
            )
            return transcription.strip()
        except Exception:
            return None

    def _process_audio_hybrid(self, audio_data):
        transcript_text = self._transcribe_groq(audio_data)
        
        if not transcript_text or len(transcript_text) <= 2: return ""

        clean_text = transcript_text.lower().strip('.').strip()
        for phrase in HALLUCINATION_PHRASES:
            if phrase.lower() in clean_text or clean_text == phrase.lower():
                return ""
        
        if re.match(r'^[\W_]+$', transcript_text): return ""
        
        words = clean_text.split()
        if len(words) > 4 and len(set(words)) == 1: 
            return ""

        try:
            tgt_lang = self.settings.get('target_lang', 'ja')
            tgt_name = "Japanese" if tgt_lang == 'ja' else "English"
            prompt = f"""
            Target Transcript: "{transcript_text}"
            Instruction:
            Translate the transcript into natural {tgt_name}.
            Do NOT include speaker labels.
            Output ONLY the translated text.
            """
            
            wav_file = self._create_wav_bytes(audio_data)
            if not wav_file: return ""
            wav_bytes = wav_file.getvalue()

            contents = [prompt, types.Part.from_bytes(data=wav_bytes, mime_type="audio/wav")]
            config = types.GenerateContentConfig(temperature=0.0, system_instruction=self.system_instruction)
            response = self.client.models.generate_content(model=self.model_name, contents=contents, config=config)
            
            raw_text = response.text.strip()
            if not raw_text: return ""
            if tgt_lang == 'ja':
                filtered = [l for l in raw_text.split('\n') if any(ord(c) > 128 for c in l.split(':', 1)[-1])]
                return "\n".join(filtered)
            return raw_text
        except Exception: 
            return ""

    def _main_loop(self):
        print(f"🔁 処理開始 [User: {self.sid}] URL: {self.url}")
        self.download_thread = threading.Thread(target=self._run_ytdlp)
        self.download_thread.daemon = True
        self.download_thread.start()

        for _ in range(200):
            if os.path.exists(self.fifo_path):
                break
            time.sleep(0.1)

        if not os.path.exists(self.fifo_path):
            self.stop()
            return

        ffmpeg_cmd = [
            "ffmpeg", "-y", "-i", self.fifo_path,
            "-f", "s16le", "-acodec", "pcm_s16le",
            "-ac", "1", "-ar", "16000", "-bufsize", "0", "pipe:1"
        ]
        try:
            self.ffmpeg_proc = subprocess.Popen(
                ffmpeg_cmd, stdout=subprocess.PIPE, stderr=None
            )
        except Exception:
            self.stop()
            return

        print(f"👂 音声待機中... ({self.sid})")

        CHUNK_DURATION = 0.1
        SAMPLE_RATE = 16000
        BYTES_PER_SAMPLE = 2
        CHUNK_BYTES = int(SAMPLE_RATE * CHUNK_DURATION * BYTES_PER_SAMPLE)

        SILENCE_THRESHOLD = 300
        SILENCE_DURATION_TRIGGER = 0.6
        MAX_RECORD_DURATION = 6.0
        MIN_RECORD_DURATION = 1.0

        audio_buffer = bytearray()
        silence_counter = 0.0
        recording_duration = 0.0

        while self.is_running:
            try:
                chunk = self.ffmpeg_proc.stdout.read(CHUNK_BYTES)
                if not chunk:
                    break

                audio_np = np.frombuffer(chunk, dtype=np.int16)
                if len(audio_np) == 0 or np.isnan(audio_np).any():
                    continue

                rms = np.sqrt(np.mean(audio_np.astype(float) ** 2))

                SPEAKER_SWITCH_THRESHOLD = 600
                speaker_switch_candidate = abs(rms - self.last_rms) > SPEAKER_SWITCH_THRESHOLD
                self.last_rms = rms

                audio_buffer.extend(chunk)
                recording_duration += CHUNK_DURATION

                if rms < SILENCE_THRESHOLD:
                    silence_counter += CHUNK_DURATION
                else:
                    silence_counter = 0.0

                should_process = (
                    (recording_duration > MIN_RECORD_DURATION and silence_counter >= SILENCE_DURATION_TRIGGER)
                    or recording_duration >= MAX_RECORD_DURATION
                )

                if not should_process:
                    continue

                if speaker_switch_candidate:
                    max_spk = self.settings.get("max_speakers", 1)
                    if max_spk > 1:
                        idx = int(self.current_speaker.split()[-1])
                        idx = (idx % max_spk) + 1
                        self.current_speaker = f"Speaker {idx}"

                audio_bytes = bytes(audio_buffer)
                result_text = self._process_audio_hybrid(audio_bytes)

                audio_buffer.clear()
                recording_duration = 0.0
                silence_counter = 0.0

                if not result_text:
                    continue

                raw_id = self.current_speaker
                text_content = result_text.strip()
                if not text_content:
                    continue

                max_spk = self.settings.get("max_speakers", 1)
                if max_spk == 1:
                    raw_id = "Speaker 1"

                spk_match = re.search(r"\d+", raw_id)
                if spk_match and int(spk_match.group()) > max_spk:
                    raw_id = "Speaker 1"

                if self.sid not in session_speakers:
                    session_speakers[self.sid] = {}

                if raw_id not in session_speakers[self.sid]:
                    new_idx = len(session_speakers[self.sid])
                    session_speakers[self.sid][raw_id] = {
                        "name": raw_id,
                        "color": SPEAKER_COLORS[new_idx % len(SPEAKER_COLORS)]
                    }
                    self.socketio.emit(
                        "initial_speaker_profiles",
                        session_speakers[self.sid],
                        room=self.sid
                    )

                info = session_speakers[self.sid][raw_id]

                self.socketio.emit("new_subtitle", {
                    "original": "...",
                    "translated": text_content,
                    "speaker": info["name"],
                    "speaker_id": raw_id,
                    "color": info["color"]
                }, room=self.sid)

            except Exception:
                continue

        self.stop()


@app.route("/")
def index(): 
    return render_template("index.html")

@socketio.on('connect')
def handle_connect(): 
    pass

@socketio.on('disconnect')
def handle_disconnect():
    if request.sid in active_sessions: 
        active_sessions[request.sid].stop()
        del active_sessions[request.sid]
    if request.sid in session_speakers: 
        del session_speakers[request.sid]

@socketio.on('start_stream')
def handle_start_stream(data):
    sid = request.sid
    if sid in active_sessions: 
        active_sessions[sid].stop()
    if len(active_sessions) >= MAX_CONCURRENT_STREAMS: 
        emit("error_msg", {"msg": "混雑中"}, room=sid)
        return
    worker = StreamWorker(
        sid, 
        data.get('url'), 
        {'target_lang': data.get('target_lang', 'ja'), 'max_speakers': int(data.get('max_speakers', 1))}, 
        socketio
    )
    active_sessions[sid] = worker
    worker.start()
    emit("status_msg", {"msg": "開始"}, room=sid)

@socketio.on('stop_stream')
def handle_stop_stream():
    if request.sid in active_sessions: 
        active_sessions[request.sid].stop()
        emit("status_msg", {"msg": "停止"}, room=request.sid)

@socketio.on('update_settings')
def handle_settings_update(data):
    if request.sid in active_sessions: 
        active_sessions[request.sid].settings.update(data)

@socketio.on('reset_speakers')
def handle_reset_speakers():
    if request.sid in session_speakers: 
        session_speakers[request.sid] = {}

@socketio.on('update_speaker_name')
def handle_speaker_name_update(data):
    if request.sid in session_speakers and data.get("id") in session_speakers[request.sid]:
        session_speakers[request.sid][data.get("id")]["name"] = data.get("new_name")
        emit("speaker_name_changed", data, room=request.sid)

@socketio.on('update_speaker_color')
def handle_speaker_color_update(data):
    if request.sid in session_speakers and data.get("id") in session_speakers[request.sid]:
        session_speakers[request.sid][data.get("id")]["color"] = data.get("color")
        emit("speaker_color_changed", data, room=request.sid)

if __name__ == "__main__":
    socketio.run(app, host='0.0.0.0', port=7860, allow_unsafe_werkzeug=True)