#!/usr/bin/env python3
"""
YouTube Factory — Master Orchestrator for 24/7 Faceless Video Generation
Distributed across 3 AI PCs:
- AMDLLM:    LLM Scriptwriting, Translation & Trending Ideation (llama.cpp Qwen3.5-7B)
- NvidiaLLM: ComfyUI MiniMax H3 Video Gen, z-image-turbo / SDXL & A/B Thumbnails
- IntelLLM:  Kokoro TTS (EN+ES), Faster-Whisper Subtitles & FFmpeg Assembly with Audio Ducking
"""

import os
import sys
import json
import yaml
import time
import re
import argparse
import subprocess
import logging
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any

from pipeline.remote import RemotePC
from pipeline.trending import TrendingFetcher
from pipeline.script_gen import ScriptGen
from pipeline.tts_gen import TTSGen
from pipeline.music_gen import MusicGen
from pipeline.visual_gen import VisualGen
from pipeline.sub_gen import SubGen
from pipeline.thumb_gen import ThumbnailGen
from pipeline.assembler import Assembler
from pipeline.queue_manager import QueueManager
from pipeline.youtube_uploader import YouTubeUploader
from pipeline.dependency_manager import ensure_dependencies

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(message)s",
    datefmt="%H:%M:%S"
)
logger = logging.getLogger("YTFactory")

@dataclass
class FactoryConfig:
    topic: str = ""
    duration_seconds: int = 180
    output_dir: str = "./output"
    temp_dir: str = "./temp"
    visual_strategy: str = "hybrid"
    clip_width: int = 768
    clip_height: int = 432
    voices: Dict[str, str] = field(default_factory=lambda: {"en": "af_heart", "es": "ef_dora"})
    thumbnail_count: int = 3
    music_mood: str = "cinematic"
    pexels_api_key: Optional[str] = None
    pixabay_api_key: Optional[str] = None
    poll_interval_seconds: int = 30
    auto_trending_when_empty: bool = True
    trending_niche: str = ""
    max_daily_videos: int = 12
    youtube_client_secrets: str = "client_secrets.json"
    youtube_token_file: str = "youtube_token.json"
    youtube_default_privacy: str = "private"
    youtube_auto_upload: bool = False

def format_timestamp(seconds: float) -> str:
    """Format seconds into MM:SS format for YouTube chapter markers."""
    m = int(seconds) // 60
    s = int(seconds) % 60
    return f"{m:02d}:{s:02d}"

class YouTubeFactory:
    def __init__(self, config_path: str = "factory_config.yaml", check_deps: bool = True):
        if check_deps:
            ensure_dependencies(auto_install=True)

        self.config_path = config_path
        self.raw_config = self._load_config(config_path)
        self.cfg = self._parse_config(self.raw_config)
        self.pcs: Dict[str, RemotePC] = {}
        for name, data in self.raw_config.get("pcs", {}).items():
            self.pcs[name] = RemotePC(name, data["host"], data["ssh_user"], data.get("services", {}))
        
        self.queue = QueueManager(db_path=os.path.join(self.cfg.output_dir, "factory_queue.db"))
        self.uploader = YouTubeUploader(
            client_secrets_file=self.cfg.youtube_client_secrets,
            token_file=self.cfg.youtube_token_file,
            default_privacy=self.cfg.youtube_default_privacy
        )
        os.makedirs(self.cfg.output_dir, exist_ok=True)
        os.makedirs(self.cfg.temp_dir, exist_ok=True)

    def _load_config(self, path: str) -> Dict[str, Any]:
        if not os.path.exists(path):
            raise FileNotFoundError(f"Config file not found: {path}")
        with open(path, "r", encoding="utf-8") as f:
            return yaml.safe_load(f)

    def _parse_config(self, raw: Dict[str, Any]) -> FactoryConfig:
        p_cfg = raw.get("pipeline", {})
        d_cfg = raw.get("daemon", {})
        y_cfg = raw.get("youtube", {})
        return FactoryConfig(
            topic=p_cfg.get("topic", ""),
            duration_seconds=p_cfg.get("duration_seconds", 180),
            output_dir=p_cfg.get("output_dir", "./output"),
            temp_dir=p_cfg.get("temp_dir", "./temp"),
            visual_strategy=p_cfg.get("visual_strategy", "hybrid"),
            clip_width=p_cfg.get("clip_width", 768),
            clip_height=p_cfg.get("clip_height", 432),
            voices=p_cfg.get("voices", {"en": "af_heart", "es": "ef_dora"}),
            thumbnail_count=p_cfg.get("thumbnail_count", 3),
            music_mood=p_cfg.get("music_mood", "cinematic"),
            pexels_api_key=p_cfg.get("pexels_api_key") or None,
            pixabay_api_key=p_cfg.get("pixabay_api_key") or None,
            poll_interval_seconds=d_cfg.get("poll_interval_seconds", 30),
            auto_trending_when_empty=d_cfg.get("auto_trending_when_empty", True),
            trending_niche=d_cfg.get("trending_niche", ""),
            max_daily_videos=d_cfg.get("max_daily_videos", 12),
            youtube_client_secrets=y_cfg.get("client_secrets_file", "client_secrets.json"),
            youtube_token_file=y_cfg.get("token_file", "youtube_token.json"),
            youtube_default_privacy=y_cfg.get("default_privacy", "private"),
            youtube_auto_upload=y_cfg.get("auto_upload", False)
        )

    def health(self):
        """Check connection and API service availability on all 3 PCs."""
        logger.info("=" * 65)
        logger.info("📊 3-PC YouTube Factory Health Check")
        logger.info("=" * 65)
        all_ok = True
        for name, pc in self.pcs.items():
            logger.info(f"Checking PC: {name.upper()} ({pc.host})...")
            for svc_name, svc_info in pc.services.items():
                port = svc_info.get("port", 0)
                if port > 0:
                    status = pc.check(port)
                    icon = "✅ UP  " if status else "❌ DOWN"
                    logger.info(f"  {icon} {name}.{svc_name:<12} on port :{port}")
                    if not status:
                        all_ok = False
        logger.info("=" * 65)
        if all_ok:
            logger.info("🎉 All AI services are responding and ready.")
        else:
            logger.warning("⚠️ Some services are currently unreachable. Review above.")

    def run_video_pipeline(self, topic: str, duration: Optional[int] = None, dry_run: bool = False) -> Dict[str, Any]:
        """Execute the end-to-end video production pipeline for a single topic."""
        target_dur = duration or self.cfg.duration_seconds
        clean_topic_slug = re.sub(r'[^A-Za-z0-9_]', '_', topic[:25]).strip('_')
        job_id = f"{int(time.time())}_{clean_topic_slug}"
        job_dir = os.path.join(self.cfg.temp_dir, job_id)
        os.makedirs(job_dir, exist_ok=True)

        logger.info("═" * 70)
        logger.info(f"🎬 NEW VIDEO JOB: \"{topic}\"")
        logger.info(f"⏱️  TARGET DURATION: {target_dur}s ({target_dur//60}m {target_dur%60}s)")
        logger.info(f"📁 JOB ID: {job_id}")
        logger.info("═" * 70)

        timings = {}
        def time_stage(name, fn):
            t0 = time.time()
            res = fn()
            elapsed = round(time.time() - t0, 1)
            timings[name] = elapsed
            logger.info(f"[{name}] Completed in {elapsed}s")
            return res

        # ── Step 1: Script Generation & Translation (AMDLLM) ──
        script_file = os.path.join(job_dir, "script.json")
        if os.path.exists(script_file):
            logger.info("[SCRIPT] Loading existing script checkpoint...")
            with open(script_file, "r", encoding="utf-8") as f:
                script = json.load(f)
        else:
            script_gen = ScriptGen(self.pcs)
            script = time_stage("SCRIPT", lambda: script_gen.generate(topic, target_dur))
            with open(script_file, "w", encoding="utf-8") as f:
                json.dump(script, f, indent=2, ensure_ascii=False)

        if dry_run:
            logger.info("[DRY-RUN] Script and metadata generated successfully. Stopping before GPU rendering.")
            return {"job_id": job_id, "script": script}

        # ── Step 2: Dual Voiceover Synthesis (IntelLLM) ──
        tts_gen = TTSGen(self.pcs, self.cfg.voices)
        tts_en = time_stage("TTS:EN", lambda: tts_gen.generate(script, job_dir, "en"))
        tts_es = time_stage("TTS:ES", lambda: tts_gen.generate(script, job_dir, "es"))

        # ── Step 3: Procedural Background Music Generation ──
        music_gen = MusicGen()
        music_path = os.path.join(job_dir, "background_music.wav")
        max_voice_dur = max(tts_en["total_duration"], tts_es["total_duration"])
        time_stage("MUSIC", lambda: music_gen.generate(max_voice_dur + 5.0, music_path, mood=self.cfg.music_mood))

        # ── Step 4: Visual Generation & Stock Media (NvidiaLLM & Web) ──
        visual_gen = VisualGen(self.pcs, self.cfg)
        clips = time_stage("VISUAL", lambda: visual_gen.generate(script, job_dir, segment_timings=tts_en["segments"]))

        # ── Step 5: Subtitle Transcription via Faster-Whisper (IntelLLM) ──
        sub_gen = SubGen(self.pcs)
        subs_en = time_stage("SUBS:EN", lambda: sub_gen.generate(tts_en["master_wav"], job_dir, "en"))
        subs_es = time_stage("SUBS:ES", lambda: sub_gen.generate(tts_es["master_wav"], job_dir, "es"))

        # ── Step 6: A/B Thumbnail Generation (NvidiaLLM) ──
        thumb_gen = ThumbnailGen(self.pcs, self.cfg)
        thumbs = time_stage("THUMBNAILS", lambda: thumb_gen.generate(script, topic, job_dir))

        # ── Step 7: Final Master Video Assembly with Audio Ducking (IntelLLM) ──
        assembler = Assembler(self.pcs)
        out_prefix = os.path.join(self.cfg.output_dir, job_id)
        rendered_videos = time_stage("ASSEMBLY", lambda: assembler.assemble(
            clips=clips,
            audio_map={"en": tts_en["master_wav"], "es": tts_es["master_wav"]},
            music_path=music_path,
            subs_map={"en": subs_en, "es": subs_es},
            script=script,
            out_prefix=out_prefix
        ))

        # ── Step 8: Package Ready-to-Upload Bundles & Copy Assets ──
        # Copy thumbnails to output folder
        packaged_thumbs = {"en": [], "es": [], "clean": []}
        for lang_key in ("en", "es", "clean"):
            for t_file in thumbs.get(lang_key, []):
                dest = f"{out_prefix}_{os.path.basename(t_file)}"
                subprocess.run(["cp", t_file, dest], check=True)
                packaged_thumbs[lang_key].append(dest)

        # Copy standalone subtitle tracks to output
        packaged_subs = {}
        for l_code, s_file in [("en", subs_en), ("es", subs_es)]:
            s_dest = f"{out_prefix}_subtitles_{l_code}.srt"
            subprocess.run(["cp", s_file, s_dest], check=True)
            packaged_subs[l_code] = s_dest

        # Build Formatted Chapters for Descriptions
        chapters_en_txt = "\n".join([
            f"{format_timestamp(s['start'])} - {s['chapter_title']}"
            for s in tts_en["segments"]
        ])
        chapters_es_txt = "\n".join([
            f"{format_timestamp(s['start'])} - {s.get('chapter_title', f'Parte {i+1}')}"
            for i, s in enumerate(tts_es["segments"])
        ])

        desc_en = f"{script.get('description', '')}\n\nTIMESTAMPS / CHAPTERS:\n{chapters_en_txt}\n\n#Documentary #History #Science"
        desc_es = f"{script.get('description_es', script.get('description', ''))}\n\nCAPÍTULOS:\n{chapters_es_txt}\n\n#Documental #Historia #Ciencia"

        # Master Metadata Manifest
        bundle_meta = {
            "job_id": job_id,
            "topic": topic,
            "titles": script.get("titles", [topic]),
            "titles_es": script.get("titles_es", [topic]),
            "description": desc_en,
            "description_es": desc_es,
            "tags": script.get("tags", []),
            "chapters_en": tts_en["segments"],
            "chapters_es": tts_es["segments"],
            "clips": clips,
            "files": rendered_videos,
            "subtitles": packaged_subs,
            "thumbnails": packaged_thumbs,
            "timings": timings,
            "created_at": time.time()
        }

        meta_json_path = f"{out_prefix}_meta.json"
        with open(meta_json_path, "w", encoding="utf-8") as f:
            json.dump(bundle_meta, f, indent=2, ensure_ascii=False)

        # Upload Text Files (Copy-paste friendly for YouTube Studio)
        with open(f"{out_prefix}_upload_EN.txt", "w", encoding="utf-8") as f:
            f.write("=" * 60 + "\n")
            f.write(f"YOUTUBE UPLOAD BUNDLE — ENGLISH (Topic: {topic})\n")
            f.write("=" * 60 + "\n\n")
            f.write("TITLES FOR A/B TESTING:\n")
            for i, t in enumerate(script.get("titles", [])):
                f.write(f"  {chr(65+i)}. {t}\n")
            f.write(f"\nDESCRIPTION & CHAPTERS:\n{desc_en}\n\n")
            f.write(f"TAGS:\n{', '.join(script.get('tags', []))}\n\n")
            f.write(f"VIDEO FILE: {rendered_videos.get('en', '')}\n")
            f.write(f"SUBTITLES:  {packaged_subs.get('en', '')}\n")
            f.write("THUMBNAILS:\n")
            for t in packaged_thumbs.get("en", []):
                f.write(f"  - {t}\n")

        with open(f"{out_prefix}_upload_ES.txt", "w", encoding="utf-8") as f:
            f.write("=" * 60 + "\n")
            f.write(f"PAQUETE DE SUBIDA A YOUTUBE — ESPAÑOL (Tema: {topic})\n")
            f.write("=" * 60 + "\n\n")
            f.write("TÍTULOS PARA PRUEBA A/B:\n")
            for i, t in enumerate(script.get("titles_es", [])):
                f.write(f"  {chr(65+i)}. {t}\n")
            f.write(f"\nDESCRIPCIÓN Y CAPÍTULOS:\n{desc_es}\n\n")
            f.write(f"ETIQUETAS / TAGS:\n{', '.join(script.get('tags', []))}\n\n")
            f.write(f"ARCHIVO DE VIDEO: {rendered_videos.get('es', '')}\n")
            f.write(f"SUBTÍTULOS:       {packaged_subs.get('es', '')}\n")
            f.write("MINIATURAS:\n")
            for t in packaged_thumbs.get("es", []):
                f.write(f"  - {t}\n")

        logger.info("═" * 70)
        logger.info("🎉 VIDEO GENERATION COMPLETE!")
        logger.info(f"📹 English Video: {rendered_videos.get('en')}")
        logger.info(f"📹 Spanish Video: {rendered_videos.get('es')}")
        logger.info(f"🖼️  Thumbnails:    {len(packaged_thumbs.get('en', []))} A/B variants")
        logger.info(f"📄 Manifest:      {meta_json_path}")
        logger.info("═" * 70)

        # Optional Auto-Upload
        if self.cfg.youtube_auto_upload and self.uploader.is_configured():
            try:
                logger.info("[YOUTUBE] Auto-upload enabled. Publishing to YouTube...")
                self.uploader.upload_bundle(meta_json_path, lang="en")
            except Exception as e:
                logger.error(f"[YOUTUBE] Auto-upload failed: {e}")

        return bundle_meta

    def run_daemon(self):
        """24/7 Autonomous Daemon Worker Loop."""
        logger.info("=" * 65)
        logger.info("🤖 YouTube Factory 24/7 Daemon Worker Started")
        logger.info(f"   Polling Interval: {self.cfg.poll_interval_seconds}s")
        logger.info(f"   Auto-Trending:    {self.cfg.auto_trending_when_empty}")
        logger.info("=" * 65)

        trending_fetcher = TrendingFetcher(self.pcs.get("amdllm"))

        while True:
            try:
                job = self.queue.get_next_job()
                if not job:
                    if self.cfg.auto_trending_when_empty:
                        logger.info("[DAEMON] Queue is empty. Fetching next trending viral topic...")
                        trending_topics = trending_fetcher.get_trending_topics(
                            niche=self.cfg.trending_niche,
                            count=1
                        )
                        if trending_topics:
                            new_topic = trending_topics[0]
                            logger.info(f"[DAEMON] Enqueuing trending topic: '{new_topic}'")
                            self.queue.add_job(new_topic)
                            continue
                    
                    logger.debug(f"[DAEMON] Waiting {self.cfg.poll_interval_seconds}s for jobs...")
                    time.sleep(self.cfg.poll_interval_seconds)
                    continue

                # Process active job
                job_id = job["job_id"]
                topic = job["topic"]
                dur = job.get("duration") or self.cfg.duration_seconds

                logger.info(f"[DAEMON] Starting Job: {job_id} -> '{topic}'")
                self.queue.update_progress(job_id, "PROCESSING")

                try:
                    result = self.run_video_pipeline(topic, duration=dur)
                    self.queue.mark_completed(job_id, result)
                except Exception as e:
                    logger.error(f"[DAEMON] Error processing job {job_id}: {e}", exc_info=True)
                    self.queue.mark_failed(job_id, str(e))

            except KeyboardInterrupt:
                logger.info("\n[DAEMON] Shutting down 24/7 worker...")
                break
            except Exception as e:
                logger.error(f"[DAEMON] Unexpected worker error: {e}", exc_info=True)
                time.sleep(10)

def main():
    parser = argparse.ArgumentParser(description="YouTube Factory — 3-PC Faceless Video Pipeline")
    parser.add_argument("--topic", "-t", help="Single video topic to generate")
    parser.add_argument("--duration", "-d", type=int, help="Target duration in seconds")
    parser.add_argument("--batch", "-b", nargs="+", help="List of video topics to process sequentially")
    parser.add_argument("--batch-file", "-f", help="Text file containing one topic per line")
    parser.add_argument("--trending", action="store_true", help="Fetch real-time trending topics and generate videos")
    parser.add_argument("--trending-niche", help="Trending niche (history, space, science, tech, mysteries, business)")
    parser.add_argument("--trending-count", type=int, default=3, help="Number of trending topics to fetch/queue")
    parser.add_argument("--daemon", action="store_true", help="Start 24/7 continuous autonomous generation worker")
    parser.add_argument("--queue-add", help="Enqueue a new topic into the background queue")
    parser.add_argument("--queue-list", action="store_true", help="Display all pending and completed jobs in queue")
    parser.add_argument("--queue-retry", action="store_true", help="Reset all failed jobs back to queued state")
    parser.add_argument("--health", "-H", action="store_true", help="Check status of all 3 PCs and their AI services")
    parser.add_argument("--setup", action="store_true", help="Check and automatically install all dependencies across environment")
    parser.add_argument("--check-deps", action="store_true", help="Verify Python packages, system tools, and SSH connections")
    parser.add_argument("--dry-run", action="store_true", help="Generate script and metadata without running heavy GPU video rendering")
    parser.add_argument("--upload", help="Upload a completed video bundle meta.json to YouTube")
    parser.add_argument("--config", "-c", default="factory_config.yaml", help="Path to custom config YAML")

    args = parser.parse_args()

    if args.setup or args.check_deps:
        print("\n🔍 Checking YouTube Factory Dependencies...\n")
        ok = ensure_dependencies(auto_install=True, check_remote=True)
        if ok:
            print("\n✅ All dependencies and remote connections are fully configured!\n")
        else:
            print("\n⚠️ Some dependencies or remote connections require attention.\n")
        sys.exit(0 if ok else 1)

    factory = YouTubeFactory(args.config)

    if args.health:
        factory.health()
    elif args.upload:
        if not factory.uploader.is_configured():
            logger.error("YouTube API credentials not configured. Please place client_secrets.json in this directory.")
            sys.exit(1)
        factory.uploader.upload_bundle(args.upload)
    elif args.queue_add:
        factory.queue.add_job(args.queue_add, duration=args.duration)
    elif args.queue_list:
        jobs = factory.queue.list_jobs()
        print("\n" + "=" * 80)
        print(f"{'JOB ID':<20} | {'STATUS':<12} | {'STEP':<12} | {'TOPIC'}")
        print("-" * 80)
        for j in jobs:
            print(f"{j['job_id']:<20} | {j['status']:<12} | {j['progress_step']:<12} | {j['topic'][:40]}")
        print("=" * 80 + "\n")
    elif args.queue_retry:
        retried = factory.queue.retry_failed()
        logger.info(f"Retried {retried} failed jobs.")
    elif args.daemon:
        factory.run_daemon()
    elif args.trending:
        fetcher = TrendingFetcher(factory.pcs.get("amdllm"))
        topics = fetcher.get_trending_topics(niche=args.trending_niche, count=args.trending_count)
        logger.info(f"Discovered {len(topics)} trending topics: {topics}")
        for t in topics:
            factory.run_video_pipeline(t, duration=args.duration, dry_run=args.dry_run)
    elif args.batch_file:
        if not os.path.exists(args.batch_file):
            logger.error(f"Batch file not found: {args.batch_file}")
            sys.exit(1)
        with open(args.batch_file, "r") as f:
            topics = [line.strip() for line in f if line.strip() and not line.startswith("#")]
        for t in topics:
            factory.run_video_pipeline(t, duration=args.duration, dry_run=args.dry_run)
    elif args.batch:
        for t in args.batch:
            factory.run_video_pipeline(t, duration=args.duration, dry_run=args.dry_run)
    elif args.topic:
        factory.run_video_pipeline(args.topic, duration=args.duration, dry_run=args.dry_run)
    else:
        # Interactive mode
        topic = input("Enter video topic: ").strip()
        if topic:
            factory.run_video_pipeline(topic, duration=args.duration, dry_run=args.dry_run)

if __name__ == "__main__":
    main()