import os
import json
import time
import shutil
import re
import logging
import tempfile
import random
from pathlib import Path
from collections import defaultdict

import requests
from flask import Flask, request, render_template, send_from_directory, jsonify
from flask_cors import CORS
from werkzeug.utils import secure_filename
import cv2
import fitz  # PyMuPDF
import pptx
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification
from dotenv import load_dotenv

# --- OpenAI ---
import openai
from openai import OpenAI

# --- Local Imports ---
try:
    from imagegen.gen_img_openai import main_generate_images
    from imagegen.generate_script import VideoScriptGenerator
    from tts.generate_audio import main_generate_audio
    from assembly.scripts.assembly_video import create_video, create_complete_srt
except ImportError as e:
    print(f"CRITICAL WARNING: Missing local module: {e}")

# ==========================================
# CONFIGURATION & SETUP
# ==========================================

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(module)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler("app_debug.log")
    ]
)

logger = logging.getLogger(__name__)

app = Flask(__name__)
CORS(app)

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
UPLOAD_FOLDER = os.path.join(BASE_DIR, "uploads")
STATIC_VIDEO_FOLDER = os.path.join(BASE_DIR, "static", "videos")
RESOURCE_DIR = os.path.join(BASE_DIR, "resources")

SCRIPT_PATH = os.path.join(RESOURCE_DIR, "scripts", "script.json")
IMAGES_PATH = os.path.join(RESOURCE_DIR, "images/")
AUDIO_PATH = os.path.join(RESOURCE_DIR, "audio/")
FONT_PATH = os.path.join(RESOURCE_DIR, "font", "font.ttf")
VIDEO_OUTPUT_PATH = os.path.join(RESOURCE_DIR, "video/")
SUBTITLE_OUTPUT_PATH = os.path.join(RESOURCE_DIR, "subtitles/")

for path in [
    UPLOAD_FOLDER,
    STATIC_VIDEO_FOLDER,
    IMAGES_PATH,
    AUDIO_PATH,
    VIDEO_OUTPUT_PATH,
    SUBTITLE_OUTPUT_PATH,
    os.path.dirname(SCRIPT_PATH)
]:
    os.makedirs(path, exist_ok=True)

app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = 50 * 1024 * 1024

# ==========================================
# API KEYS
# ==========================================

load_dotenv()
OPENAI_API_KEY = os.getenv("openai_api_key")

if not OPENAI_API_KEY:
    logger.warning("OpenAI API Key missing!")

client = OpenAI(api_key=OPENAI_API_KEY)

# ==========================================
# MODELS
# ==========================================

try:
    logger.info("Loading transformer models...")
    processor = AutoImageProcessor.from_pretrained(
        "DunnBC22/vit-base-patch16-224-in21k_Human_Activity_Recognition"
    )
    model_hf = AutoModelForImageClassification.from_pretrained(
        "DunnBC22/vit-base-patch16-224-in21k_Human_Activity_Recognition"
    )
    generator = VideoScriptGenerator(api_key=OPENAI_API_KEY)
    logger.info("Models loaded.")
except Exception as e:
    logger.error(f"Model init error: {e}")

# ==========================================
# HELPERS
# ==========================================

def allowed_file(filename):
    return "." in filename and filename.rsplit(".", 1)[1].lower() in {"pdf", "pptx"}

def extract_json_from_text(text):
    try:
        match = re.search(r"\{.*\}", text, re.DOTALL)
        return json.loads(match.group()) if match else json.loads(text)
    except Exception:
        return None

def ask_openai(prompt):
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.7,
    )
    return response.choices[0].message.content

def extract_text_from_pdf(pdf_path, max_chars=2000):
    text = ""
    with fitz.open(pdf_path) as doc:
        for page in doc:
            text += page.get_text()
            if len(text) > max_chars:
                break
    return text[:max_chars]

def extract_text_from_ppt(ppt_path, max_chars=2000):
    text = ""
    prs = pptx.Presentation(ppt_path)
    for slide in prs.slides:
        for shape in slide.shapes:
            if hasattr(shape, "text"):
                text += shape.text + "\n"
                if len(text) > max_chars:
                    break
    return text[:max_chars]

def clean_directory(folder_path):
    for f in os.listdir(folder_path):
        p = os.path.join(folder_path, f)
        try:
            if os.path.isfile(p):
                os.unlink(p)
            else:
                shutil.rmtree(p)
        except Exception:
            pass

# ==========================================
# ROUTES
# ==========================================

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

@app.route("/upload", methods=["POST"])
def upload_file():
    if "file" not in request.files:
        return jsonify({"error": "No file"}), 400

    file = request.files["file"]

    if file.filename == "" or not allowed_file(file.filename):
        return jsonify({"error": "Invalid file"}), 400

    try:
        filename = secure_filename(file.filename)
        path = os.path.join(app.config["UPLOAD_FOLDER"], filename)
        file.save(path)

        text = extract_text_from_pdf(path) if filename.endswith(".pdf") \
               else extract_text_from_ppt(path)

        prompt = f"""
        Generate 5 Yes/No questions from this text.
        Return JSON only:
        {{ "questions": ["..."] }}

        Text:
        {text}
        """

        raw = ask_openai(prompt)
        data = extract_json_from_text(raw)

        return jsonify({"questions": data["questions"]})

    except Exception as e:
        logger.error(e)
        return jsonify({"error": str(e)}), 500

@app.route("/videogen", methods=["POST"])
def video_gen():

    try:
        data = request.get_json()
        topic = data.get('topic')
        duration = int(data.get('duration', 60))
        key_points = data.get('key_points', [])

        clean_directory(IMAGES_PATH)
        clean_directory(AUDIO_PATH)

        script = generator.generate_script(topic, duration, key_points)
        generator.save_script(script, SCRIPT_PATH)

        main_generate_images(SCRIPT_PATH, IMAGES_PATH)
        main_generate_audio(SCRIPT_PATH, AUDIO_PATH)

        clean_topic = re.sub(r"[^A-Za-z0-9]", "_", topic)[:30]
        video_filename = f"{clean_topic}_{int(time.time())}.mp4"

        temp_video_path = os.path.join(VIDEO_OUTPUT_PATH, video_filename)
        srt_path = os.path.join(SUBTITLE_OUTPUT_PATH, f"{clean_topic}.srt")

        create_complete_srt(SCRIPT_PATH, AUDIO_PATH, srt_path, chunk_size=10)
        create_video(IMAGES_PATH, AUDIO_PATH, SCRIPT_PATH, FONT_PATH, temp_video_path)

        shutil.copy(temp_video_path, os.path.join(STATIC_VIDEO_FOLDER, video_filename))

        return jsonify({'success': True, 'videoPath': video_filename})

    except Exception as e:
        logger.error(e)
        return jsonify({'success': False, 'error': str(e)}), 500


@app.route('/api/videos')

def get_videos():

    VIDEO_DIR="./static/videos"

    videos = []

    # Ensure the video directory exists

    if os.path.exists(VIDEO_DIR):

        # Get all files in the directory

        for filename in os.listdir(VIDEO_DIR):

            # Check for common video extensions

            if filename.lower().endswith(('.mp4', '.webm', '.ogg', '.mov', '.avi')):

                videos.append({

                    'name': filename,

                    'path': os.path.join(VIDEO_DIR, filename).replace('\\', '/')

                })

    return jsonify(videos)

@app.route('/videos/<video_name>')
def serve_video(video_name):

    return send_from_directory("./static/videos", video_name)

# ==========================================
# MAIN
# ==========================================

if __name__ == "__main__":
    logger.info("Starting Flask Server...")
    app.run(host='0.0.0.0', port=5000)