#!/usr/bin/env python3
import os
import sys
import time
import json
import re
import sqlite3
import urllib.request
import urllib.error
from pathlib import Path

# ==================== CONFIGURATION ====================
# 1. Directory to watch for new audio files
# By default, this is the macOS Voice Memos (Ghi âm) folder.
WATCH_DIR = "/Users/keira/Library/Group Containers/group.com.apple.VoiceMemos.shared/Recordings"

# 2. Obsidian Vault folder where the processed Markdown files will be saved
OBSIDIAN_VAULT_DIR = "/Users/keira/Documents/Obsidian Vault Mac/TypeWhisper"

# 3. Local TypeWhisper API endpoint
API_URL = "http://127.0.0.1:8001"

# 4. Active Workflow Name in TypeWhisper App
# The script will read this workflow's exact system prompt and template from the app database.
WORKFLOW_NAME = "Advanced-Meeting Note"

# Default models (will be overridden by active TypeWhisper settings if available)
DEFAULT_ASR_MODEL = "local-zipformer-vi-68m-punc-diar"
DEFAULT_LLM_MODEL = "gemini-2.5-flash"

# Gemini API Key (fallback if api_key.txt is not found)
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
# =======================================================

SCRIPT_DIR = Path(__file__).parent.absolute()

def log(msg):
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] {msg}")

def send_notification(title, message):
    """Sends a native macOS notification"""
    try:
        # Escape double quotes for AppleScript
        msg_escaped = message.replace('"', '\\"')
        title_escaped = title.replace('"', '\\"')
        os.system(f'osascript -e \'display notification "{msg_escaped}" with title "{title_escaped}"\'')
    except Exception:
        pass

def load_api_key():
    """Loads API key from api_key.txt or fallback config"""
    key_path = SCRIPT_DIR / "api_key.txt"
    if key_path.exists():
        try:
            with open(key_path, "r", encoding="utf-8") as f:
                key = f.read().strip()
                if key:
                    # Ensure it has Bearer prefix if it doesn't already
                    if not key.lower().startswith("bearer "):
                        key = f"Bearer {key}"
                    return key
        except Exception as e:
            log(f"Failed to read api_key.txt: {e}")
            
    if GEMINI_API_KEY:
        if not GEMINI_API_KEY.lower().startswith("bearer "):
            return f"Bearer {GEMINI_API_KEY}"
        return GEMINI_API_KEY
        
    return None

def load_active_models():
    """Loads ASR and LLM models from last_used_models.json"""
    asr_model = DEFAULT_ASR_MODEL
    llm_model = DEFAULT_LLM_MODEL
    
    config_path = SCRIPT_DIR / "last_used_models.json"
    if config_path.exists():
        try:
            with open(config_path, "r", encoding="utf-8") as f:
                data = json.load(f)
                if "asr" in data:
                    asr_model = data["asr"]
                if "llm" in data:
                    llm_model = data["llm"]
            log(f"Detected active models from TypeWhisper: ASR='{asr_model}', LLM='{llm_model}'")
        except Exception as e:
            log(f"Failed to read last_used_models.json: {e}")
    else:
        log(f"No active settings found. Using defaults: ASR='{asr_model}', LLM='{llm_model}'")
            
    return asr_model, llm_model

def load_workflow_prompt(workflow_name):
    """Loads system instruction and template from the TypeWhisper app database"""
    db_path = "/Users/keira/Library/Application Support/TypeWhisper/workflows.store"
    if not os.path.exists(db_path):
        log(f"Warning: Workflows database not found at {db_path}. Using fallback prompt.")
        return None, None
        
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        cursor.execute(
            "SELECT ZBEHAVIORDATA FROM ZWORKFLOW WHERE ZNAME = ?",
            (workflow_name,)
        )
        row = cursor.fetchone()
        if row:
            behavior_data = json.loads(row[0].decode('utf-8'))
            system_instruction = behavior_data.get("settings", {}).get("instruction", "")
            user_template = behavior_data.get("fineTuning", "")
            log(f"Loaded prompts for workflow '{workflow_name}' from TypeWhisper app database.")
            return system_instruction, user_template
    except Exception as e:
        log(f"Failed to read workflows.store: {e}")
        
    return None, None

def get_file_to_process():
    """Returns the file to process based on command line arguments or the newest file"""
    if len(sys.argv) > 1:
        arg = sys.argv[1]
        
        # Check if it's an absolute path
        path = Path(arg)
        if path.is_file():
            return path
            
        # Check if it's a file name in WATCH_DIR
        path_in_watch = Path(WATCH_DIR) / arg
        if path_in_watch.is_file():
            return path_in_watch
            
        log(f"Error: Specified file '{arg}' not found.")
        sys.exit(1)
        
    # Find the newest file
    if not os.path.exists(WATCH_DIR):
        log(f"Error: Recordings folder '{WATCH_DIR}' does not exist.")
        sys.exit(1)
        
    audio_files = []
    for entry in os.scandir(WATCH_DIR):
        if entry.is_file() and entry.name.lower().endswith(('.m4a', '.wav', '.mp3', '.caf')):
            audio_files.append((entry.path, entry.stat().st_mtime))
            
    if not audio_files:
        log("No audio files found in the recordings folder.")
        sys.exit(1)
        
    # Sort by mtime descending (newest first)
    audio_files.sort(key=lambda x: x[1], reverse=True)
    newest_path = Path(audio_files[0][0])
    return newest_path

def wait_for_file_ready(filepath, stable_seconds=2):
    """Waits until the file size is stable (finished recording/writing)"""
    log(f"Target file: {filepath.name}. Checking file status...")
    last_size = -1
    for _ in range(15):  # Max 30 seconds
        try:
            current_size = os.path.getsize(filepath)
            if current_size == last_size and current_size > 0:
                log(f"File {filepath.name} is ready ({current_size / (1024*1024):.2f} MB).")
                return True
            last_size = current_size
        except Exception:
            pass
        time.sleep(stable_seconds)
    return True

def query_asr(filepath, asr_model):
    """Sends the audio file to TypeWhisper ASR endpoint"""
    log(f"Transcribing {filepath.name} using ASR model: {asr_model}...")
    url = f"{API_URL}/v1/audio/transcriptions"
    
    boundary = "----WebKitFormBoundaryTypeWhisperAuto"
    
    part_headers = (
        f"--{boundary}\r\n"
        f'Content-Disposition: form-data; name="file"; filename="{filepath.name}"\r\n'
        f'Content-Type: audio/octet-stream\r\n\r\n'
    ).encode('utf-8')
    
    part_model = (
        f"\r\n--{boundary}\r\n"
        f'Content-Disposition: form-data; name="model"\r\n\r\n'
        f"{asr_model}\r\n"
        f"--{boundary}--\r\n"
    ).encode('utf-8')
    
    try:
        with open(filepath, "rb") as f:
            file_content = f.read()
            
        body = part_headers + file_content + part_model
        
        req = urllib.request.Request(url, data=body)
        req.add_header("Content-Type", f"multipart/form-data; boundary={boundary}")
        
        with urllib.request.urlopen(req) as response:
            result = json.loads(response.read().decode('utf-8'))
            return result
    except Exception as e:
        log(f"ASR request failed: {e}")
        return None

def query_llm(raw_text, llm_model):
    """Sends the transcript to the LLM proxy for cleaning and summarization"""
    log(f"Refining transcript using LLM model: {llm_model}...")
    url = f"{API_URL}/v1/chat/completions"
    
    # 1. Try to load the active workflow prompts from the TypeWhisper database
    system_prompt, user_template = load_workflow_prompt(WORKFLOW_NAME)
    
    if system_prompt or user_template:
        # Construct messages using the exact app prompts
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        
        user_content = f"{user_template}\n\n--- TRANSCRIPT CONTENT ---\n{raw_text}"
        messages.append({"role": "user", "content": user_content})
    else:
        # Fallback to default cleaning prompt
        log("Using fallback cleaning prompt.")
        fallback_prompt = (
            "Bạn là một trợ lý AI chuyên nghiệp. Dưới đây là bản ghi chép cuộc thoại (transcript) tiếng Việt thu được từ mô hình ASR "
            "(có thể có lỗi chính tả, từ ngữ lặp, từ đệm hoặc lỗi nhận diện người nói).\n\n"
            "Nhiệm vụ của bạn là:\n"
            "1. **Clean Text**: Sửa lỗi chính tả, bỏ từ đệm (à, ừ, thì, mà, là...), chỉnh sửa câu cú cho trôi chảy tự nhiên nhưng KHÔNG làm mất đi ý nghĩa gốc và phong cách của người nói. Giữ nguyên nhãn người nói (ví dụ: `Speaker 0:`, `Speaker 1:`) ở đầu dòng.\n"
            "2. **Structure**: Chia đoạn rõ ràng, thêm tiêu đề (headings) cho các phần nội dung chính để dễ theo dõi.\n"
            "3. **Summarize**: Tạo một bản tóm tắt ngắn gọn ở đầu trang.\n\n"
            f"--- NỘI DUNG TRANSCRIPT ---\n{raw_text}"
        )
        messages = [{"role": "user", "content": fallback_prompt}]
        
    payload = {
        "model": llm_model,
        "messages": messages,
        "temperature": 0.3
    }
    
    req = urllib.request.Request(url, data=json.dumps(payload).encode('utf-8'))
    req.add_header("Content-Type", "application/json")
    
    # Load and attach authorization API Key
    api_key = load_api_key()
    if api_key:
        req.add_header("Authorization", api_key)
        
    try:
        with urllib.request.urlopen(req) as response:
            result = json.loads(response.read().decode('utf-8'))
            return result["choices"][0]["message"]["content"]
    except Exception as e:
        log(f"LLM request failed: {e}")
        return None

def main():
    try:
        # 1. Determine which file to process
        filepath = get_file_to_process()
        log(f"=== Starting TypeWhisper Transcription Workflow ===")
        log(f"Processing file: {filepath.resolve()}")
        
        # 2. Load active models from TypeWhisper app settings
        asr_model, llm_model = load_active_models()
        
        # 3. Wait for file to be fully written
        if not wait_for_file_ready(filepath):
            return
            
        send_notification("TypeWhisper Auto", f"Đang xử lý: {filepath.name}")
        
        # 4. Transcribe
        asr_result = query_asr(filepath, asr_model)
        if not asr_result:
            send_notification("TypeWhisper Auto", "❌ Lỗi: Không thể nhận dạng giọng nói.")
            return
            
        raw_text = asr_result.get("text", "")
        segments = asr_result.get("segments", [])
        
        if not raw_text.strip():
            log("ASR returned empty text. Exiting.")
            send_notification("TypeWhisper Auto", "⚠️ Bản ghi âm không có nội dung lời nói.")
            return
            
        # Check if we have speaker diarization in segments
        formatted_segments = []
        for seg in segments:
            spk = seg.get("speaker") or seg.get("speaker_id")
            text = seg.get("text", "").strip()
            if spk is not None:
                formatted_segments.append(f"Speaker {spk}: {text}")
            else:
                formatted_segments.append(text)
                
        text_to_clean = "\n".join(formatted_segments) if formatted_segments else raw_text
        
        # 5. Clean and Summarize
        llm_result = query_llm(text_to_clean, llm_model)
        if not llm_result:
            log("LLM workflow failed. Saving raw transcript instead.")
            llm_result = f"# Raw Transcript\n\n{text_to_clean}"
            
        # 6. Save to Obsidian
        os.makedirs(OBSIDIAN_VAULT_DIR, exist_ok=True)
        
        # Generate clean filename based on date and original name
        file_stat = os.stat(filepath)
        created_time = time.strftime("%Y-%m-%d_%H%M%S", time.localtime(file_stat.st_mtime))
        
        clean_name = re.sub(r'[^\w\s-]', '', filepath.stem).strip()
        clean_name = re.sub(r'\s+', '_', clean_name)
        
        obsidian_filename = f"Recording_{created_time}_{clean_name}.md"
        obsidian_path = Path(OBSIDIAN_VAULT_DIR) / obsidian_filename
        
        # Prep frontmatter
        frontmatter = (
            "---\n"
            f"type: recording\n"
            f"date: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(file_stat.st_mtime))}\n"
            f"original_file: {filepath.name}\n"
            f"asr_model: {asr_model}\n"
            f"llm_model: {llm_model}\n"
            f"workflow: {WORKFLOW_NAME}\n"
            "---\n\n"
        )
        
        with open(obsidian_path, "w", encoding="utf-8") as f:
            f.write(frontmatter + llm_result)
        log(f"Successfully saved to Obsidian: {obsidian_path.resolve()}")
        send_notification("TypeWhisper Auto", "✅ Hoàn tất! Đã lưu bài viết vào Obsidian.")
    except PermissionError as e:
        log(f"Permission Error: {e}")
        send_notification("TypeWhisper Auto", "⚠️ Lỗi quyền: Vui lòng cấp quyền Full Disk Access cho Terminal/Shortcuts trong Cài đặt hệ thống.")
    except Exception as e:
        log(f"Error: {e}")
        send_notification("TypeWhisper Auto", f"❌ Lỗi: {str(e)}")

if __name__ == "__main__":
    main()