from flask import Flask, render_template, request, jsonify, redirect, url_for, session, flash
from flask_sqlalchemy import SQLAlchemy
from flask_login import LoginManager, login_user, logout_user, login_required, current_user, UserMixin
import requests
import os
import threading
import time
import subprocess
import random
import string
from datetime import datetime
import logging
from urllib.parse import urlsplit, urlunsplit
from bs4 import BeautifulSoup
from urllib.request import Request, urlopen
import json
import re
from base64 import b64decode
from seleniumbase import SB
import psutil
from zipfile import ZipFile
from io import BytesIO
from celery import Celery
from celery.result import AsyncResult
import redis

app = Flask(__name__)
app.secret_key = os.environ.get('SECRET_KEY', 'supergeheim')

# Database Configuration
app.config['SQLALCHEMY_DATABASE_URI'] = os.environ.get('DATABASE_URL', 'sqlite:///aniworld_downloader.db')
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False

# Celery Configuration
app.config['CELERY'] = {
    'broker_url': os.environ.get('CELERY_BROKER_URL', 'redis://localhost:6379/0'),
    'result_backend': os.environ.get('CELERY_RESULT_BACKEND', 'redis://localhost:6379/0'),
    'task_ignore_result': False,
    'task_serializer': 'json',
    'result_serializer': 'json',
    'accept_content': ['json'],
    'timezone': 'UTC',
    'enable_utc': True,
    'task_soft_time_limit': 3600,  # 1 hour
    'task_time_limit': 3900,  # 1 hour 5 minutes
    'worker_prefetch_multiplier': 1,
    'task_acks_late': True,
    'worker_max_tasks_per_child': 50,
}

# Jellyfin Configuration
JELLYFIN_URL = os.environ.get('JELLYFIN_URL', 'http://dein-jellyfin-server:8096')
JELLYFIN_API_KEY = os.environ.get('JELLYFIN_API_KEY')

# VPN Configuration
VPN_CONFIG_PATH = os.environ.get('VPN_CONFIG_PATH', '/etc/openvpn/config.ovpn')
VPN_PID_FILE = "/tmp/aniworld_vpn.pid"

db = SQLAlchemy(app)
login_manager = LoginManager(app)
login_manager.login_view = 'login'

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


# Celery Setup
def make_celery(app):
    celery = Celery(
        app.import_name,
        backend=app.config['CELERY']['result_backend'],
        broker=app.config['CELERY']['broker_url']
    )
    celery.conf.update(app.config['CELERY'])

    class ContextTask(celery.Task):
        def __call__(self, *args, **kwargs):
            with app.app_context():
                return self.run(*args, **kwargs)

    celery.Task = ContextTask
    return celery


celery = make_celery(app)

# Provider and language configurations
provider_priority = ["Vidmoly", "VOE", "SpeedFiles", "Vidoza", "Doodstream", "Streamtape"]
languages = ["German", "Ger-Sub", "Eng-Sub", "English"]

# Compilation patterns for content extraction
VOE_PATTERNS = [re.compile(r"'hls': '(?P<url>.+)'"),
                re.compile(r'prompt\("Node",\s*"(?P<url>[^"]+)"'),
                re.compile(r"window\.location\.href = '(?P<url>[^']+)'")]
STREAMTAPE_PATTERN = re.compile(r'get_video\?id=[^&\'\s]+&expires=[^&\'\s]+&ip=[^&\'\s]+&token=[^&\'\s]+\'')
DOODSTREAM_PATTERN = re.compile(r"/pass_md5/[\w-]+/(?P<token>[\w-]+)")
VIDMOLY_PATTERN = re.compile(r"sources: \[{file:\"(?P<url>.*?)\"}]")
SPEEDFILES_PATTERN = re.compile(r"var _0x5opu234 = \"(?P<content>.*?)\";")


# Database Models
class User(UserMixin, db.Model):
    id = db.Column(db.Integer, primary_key=True)
    jellyfin_username = db.Column(db.String(80), unique=True, nullable=False)
    jellyfin_user_id = db.Column(db.String(100), nullable=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    downloads = db.relationship('Download', backref='user', lazy=True)


class Download(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    title = db.Column(db.String(200), nullable=False)
    season = db.Column(db.Integer)
    episode = db.Column(db.Integer)
    provider = db.Column(db.String(50), nullable=False)
    language = db.Column(db.String(20), nullable=False)
    status = db.Column(db.String(20), default='queued')  # queued, downloading, completed, failed, cancelled
    progress = db.Column(db.Float, default=0.0)
    file_path = db.Column(db.String(500))
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    user_id = db.Column(db.Integer, db.ForeignKey('user.id'), nullable=False)
    download_url = db.Column(db.String(1000))
    redirect_link = db.Column(db.String(1000))
    celery_task_id = db.Column(db.String(100))  # Celery Task ID
    error_message = db.Column(db.Text)


class ProviderError(Exception):
    def __init__(self, *args: object) -> None:
        super().__init__(*args)


class LanguageError(Exception):
    def __init__(self, *args: object) -> None:
        super().__init__(*args)


@login_manager.user_loader
def load_user(user_id):
    return User.query.get(int(user_id))


# Jellyfin Authentication Functions
def authenticate_jellyfin(username, password):
    """Authenticate user with Jellyfin server"""
    try:
        auth_url = f"{JELLYFIN_URL}/Users/authenticatebyname"
        headers = {
            'Content-Type': 'application/json',
            'X-Emby-Authorization': f'MediaBrowser Client="AniWorld Downloader", Device="WebApp", DeviceId="webapp-{random.randint(1000, 9999)}", Version="1.0.0"'
        }

        payload = {
            "Username": username,
            "Pw": password
        }

        response = requests.post(auth_url, json=payload, headers=headers)

        if response.status_code == 200:
            data = response.json()
            return {
                'user_id': data['User']['Id'],
                'username': data['User']['Name'],
                'access_token': data['AccessToken']
            }
        else:
            logger.error(f"Jellyfin authentication failed: {response.status_code}")
            return None

    except Exception as e:
        logger.error(f"Error during Jellyfin authentication: {e}")
        return None


def verify_jellyfin_token(access_token):
    """Verify if Jellyfin token is still valid"""
    try:
        headers = {
            'X-Emby-Token': access_token
        }
        response = requests.get(f"{JELLYFIN_URL}/System/Info", headers=headers)
        return response.status_code == 200
    except Exception as e:
        logger.error(f"Error verifying Jellyfin token: {e}")
        return False


# VPN Functions (Verbessert)
def get_vpn_process():
    """Find the specific VPN process started by this script"""
    try:
        if os.path.exists(VPN_PID_FILE):
            with open(VPN_PID_FILE, 'r') as f:
                pid = int(f.read().strip())

            if psutil.pid_exists(pid):
                proc = psutil.Process(pid)
                if 'openvpn' in proc.name().lower():
                    return proc

        # Fallback: Look for openvpn process with our config file
        for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
            try:
                if ('openvpn' in proc.info['name'].lower() and
                        any(VPN_CONFIG_PATH in cmd for cmd in proc.info['cmdline'])):
                    return proc
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue

    except Exception as e:
        logger.error(f"Error finding VPN process: {e}")

    return None


def disconnect_vpn():
    """Disconnect only our specific VPN connection"""
    try:
        vpn_proc = get_vpn_process()
        if vpn_proc:
            logger.info(f"Terminating VPN process with PID: {vpn_proc.pid}")
            vpn_proc.terminate()

            try:
                vpn_proc.wait(timeout=5)
            except psutil.TimeoutExpired:
                logger.warning("VPN process didn't terminate gracefully, forcing kill")
                vpn_proc.kill()

            if os.path.exists(VPN_PID_FILE):
                os.remove(VPN_PID_FILE)

            logger.info("VPN disconnected successfully")
            return True
        else:
            logger.info("No VPN process found to disconnect")
            return True

    except Exception as e:
        logger.error(f"Error disconnecting VPN: {e}")
        return False


def connect_vpn():
    """Connect to VPN and save PID"""
    try:
        proc = subprocess.Popen([
            'sudo', 'openvpn',
            '--config', VPN_CONFIG_PATH,
            '--daemon',
            '--writepid', VPN_PID_FILE
        ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

        time.sleep(2)

        if os.path.exists(VPN_PID_FILE):
            with open(VPN_PID_FILE, 'r') as f:
                pid = int(f.read().strip())

            if psutil.pid_exists(pid):
                logger.info(f"VPN connected successfully with PID: {pid}")
                return True

        logger.error("VPN connection failed - PID file not created")
        return False

    except Exception as e:
        logger.error(f"Error connecting VPN: {e}")
        return False


def reconnect_vpn():
    """Reconnect to VPN without affecting other VPN connections"""
    try:
        logger.info("Starting VPN reconnection...")

        if not disconnect_vpn():
            logger.error("Failed to disconnect VPN")
            return False

        time.sleep(3)

        if connect_vpn():
            time.sleep(5)
            logger.info("VPN reconnected successfully")
            return True
        else:
            logger.error("Failed to reconnect VPN")
            return False

    except Exception as e:
        logger.error(f"Error reconnecting VPN: {e}")
        return False


# Content extraction functions (keeping existing ones)
def get_titles_dict(url):
    """Get titles dictionary from streaming sites"""
    try:
        html_response = urlopen(url)
        soup = BeautifulSoup(html_response, "html.parser")
        matching_li_elements = str(soup.find_all("li")).split("</a></li>, <li><a ")
        index_list = []
        for i in matching_li_elements:
            if "data-alternative-title" not in i:
                index_list.append(matching_li_elements.index(i))
        index_list = sorted(index_list, reverse=True)
        for i in index_list:
            matching_li_elements.pop(i)
        title_link_dict = {}
        for i in range(len(matching_li_elements)):
            matching_li_elements[i] = matching_li_elements[i][24:]
        for i in matching_li_elements:
            title_link_dict.update({i[i.index('" title') + 9:i.index(" Stream anschauen")] + "," + i[:i.index('"')]: i[
                                                                                                                     i.index(
                                                                                                                         'href="') + 20: i.index(
                                                                                                                         '" title')]})
        return title_link_dict
    except Exception as e:
        logger.error(f"Error getting titles from {url}: {e}")
        return {}


# [Alle anderen Content-Extraction-Funktionen bleiben gleich...]
# (get_titles_dict_2, restructure_dict, extract_lang_key_mapping, etc.)

# CELERY TASKS (Getrennte Download-Logik)
@celery.task(bind=True)
def download_episode_task(self, download_id):
    """Celery task for downloading episodes"""
    try:
        # Update task ID in database
        download = Download.query.get(download_id)
        if not download:
            logger.error(f"Download {download_id} not found")
            return {'status': 'failed', 'error': 'Download not found'}

        download.celery_task_id = self.request.id
        download.status = 'downloading'
        db.session.commit()

        # Get or find content URL
        if not download.download_url:
            if not download.redirect_link:
                error_msg = f"No redirect link for download {download_id}"
                logger.error(error_msg)
                download.status = 'failed'
                download.error_message = error_msg
                db.session.commit()
                return {'status': 'failed', 'error': error_msg}

            content_url = find_content_url(download.redirect_link, download.provider)
            if not content_url:
                error_msg = f"Could not find content URL for {download.redirect_link}"
                logger.error(error_msg)
                download.status = 'failed'
                download.error_message = error_msg
                db.session.commit()
                return {'status': 'failed', 'error': error_msg}

            download.download_url = content_url
            db.session.commit()

        # Create file path
        base_dir = os.path.join(app.root_path, "downloads", str(download.user_id), download.title)
        if download.season:
            file_path = os.path.join(base_dir, f"Season {download.season}",
                                     f"S{download.season}-E{download.episode}-{download.title}.mp4")
        else:
            file_path = os.path.join(base_dir, "Movies", f"Movie {download.episode}-{download.title}.mp4")

        # Create directory structure
        os.makedirs(os.path.dirname(file_path), exist_ok=True)

        # Download with ffmpeg
        ffmpeg_cmd = ["ffmpeg", "-i", download.download_url, "-c", "copy", "-nostdin", file_path]

        if download.provider == "Doodstream":
            ffmpeg_cmd.insert(1, "Referer: https://d0000d.com/")
            ffmpeg_cmd.insert(1, "-headers")
        elif download.provider == "Vidmoly":
            ffmpeg_cmd.insert(1, "Referer: https://vidmoly.to/")
            ffmpeg_cmd.insert(1, "-headers")

        process = subprocess.Popen(ffmpeg_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        # Monitor download progress
        start_time = time.time()
        while process.poll() is None:
            # Update progress
            elapsed = time.time() - start_time
            if elapsed > 0:
                progress = min(95.0, (elapsed / 300) * 100)  # Estimate 5 minutes
                download.progress = progress
                db.session.commit()

                # Update Celery task progress
                self.update_state(
                    state='PROGRESS',
                    meta={'current': progress, 'total': 100, 'status': f'Downloading {download.title}...'}
                )

            time.sleep(3)

        # Check if download completed successfully
        if process.returncode == 0 and os.path.exists(file_path):
            download.status = 'completed'
            download.progress = 100.0
            download.file_path = file_path
            db.session.commit()

            logger.info(f"Download completed: {file_path}")

            # Reconnect VPN after download
            reconnect_vpn()

            return {'status': 'completed', 'file_path': file_path}
        else:
            error_msg = f"Download failed: {file_path}"
            download.status = 'failed'
            download.error_message = error_msg
            db.session.commit()

            logger.error(error_msg)

            # Clean up partial file
            if os.path.exists(file_path):
                os.remove(file_path)

            return {'status': 'failed', 'error': error_msg}

    except Exception as e:
        error_msg = f"Error downloading episode {download_id}: {str(e)}"
        logger.error(error_msg)

        download = Download.query.get(download_id)
        if download:
            download.status = 'failed'
            download.error_message = error_msg
            db.session.commit()

        return {'status': 'failed', 'error': error_msg}


# [Alle anderen Hilfsfunktionen bleiben gleich...]
# (find_content_url, get_season_count, get_episodes, etc.)

# Routes
@app.route('/')
@login_required
def index():
    return render_template('index.html')


@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        if request.is_json:
            data = request.get_json()
            username = data.get('username')
            password = data.get('password')
        else:
            username = request.form['username']
            password = request.form['password']

        jellyfin_user = authenticate_jellyfin(username, password)
        if jellyfin_user:
            user = User.query.filter_by(jellyfin_username=username).first()
            if not user:
                user = User(
                    jellyfin_username=username,
                    jellyfin_user_id=jellyfin_user['user_id']
                )
                db.session.add(user)
                db.session.commit()
            login_user(user)
            if request.is_json:
                return jsonify({'success': True})
            return redirect(url_for('index'))
        else:
            if request.is_json:
                return jsonify({'success': False, 'error': 'Login fehlgeschlagen'}), 401
            flash('Ungültige Jellyfin-Anmeldedaten')

    return render_template('login.html')


@app.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('login'))


@app.route('/api/download/start', methods=['POST'])
@login_required
def start_download():
    data = request.json
    link = data.get('link')
    provider = data.get('provider', 'Vidmoly')
    language = data.get('language', 'German')
    content_type = data.get('content_type', 'episodes')

    if not link:
        return jsonify({'error': 'No link provided'}), 400

    normalized_link = validate_link(link)
    if not normalized_link:
        return jsonify({'error': 'Invalid link format'}), 400

    title = extract_title(normalized_link)
    site = normalized_link[:normalized_link.index(".to/") + 3]

    download_ids = []
    task_ids = []

    if content_type == 'episodes':
        season_start = data.get('season_start', 1)
        season_end = data.get('season_end', 1)
        episode_start = data.get('episode_start', 1)
        episode_end = data.get('episode_end', 1)

        for season in range(season_start, season_end + 1):
            ep_start = episode_start if season == season_start else 1
            ep_end = episode_end if season == season_end else get_episodes(normalized_link, season)

            for episode in range(ep_start, ep_end + 1):
                if "bs.to" in normalized_link:
                    episode_link = normalized_link + "{}/{}-".format(season, episode)
                else:
                    episode_link = normalized_link + "staffel-{}/episode-{}".format(season, episode)

                try:
                    redirect_link = get_redirect_link_by_provider(
                        site, episode_link, language, provider, season, episode, current_user.id
                    )

                    if redirect_link:
                        download = Download(
                            title=title,
                            season=season,
                            episode=episode,
                            provider=provider,
                            language=language,
                            redirect_link=redirect_link,
                            user_id=current_user.id
                        )
                        db.session.add(download)
                        db.session.commit()
                        download_ids.append(download.id)

                        # Start Celery task
                        task = download_episode_task.delay(download.id)
                        task_ids.append(task.id)

                        # Update download with task ID
                        download.celery_task_id = task.id
                        db.session.commit()

                except Exception as e:
                    logger.error(f"Error adding download for S{season}E{episode}: {e}")

    return jsonify({
        'success': True,
        'downloads_added': len(download_ids),
        'download_ids': download_ids,
        'task_ids': task_ids
    })


@app.route('/api/download/<int:download_id>/status')
@login_required
def download_status(download_id):
    download = Download.query.get_or_404(download_id)
    if download.user_id != current_user.id:
        return jsonify({'error': 'Unauthorized'}), 403

    result_data = {
        'id': download.id,
        'status': download.status,
        'progress': download.progress,
        'error_message': download.error_message
    }

    # Get Celery task status if available
    if download.celery_task_id:
        try:
            task_result = AsyncResult(download.celery_task_id, app=celery)
            if task_result.state == 'PROGRESS':
                result_data['celery_progress'] = task_result.info
            elif task_result.state == 'SUCCESS':
                result_data['celery_result'] = task_result.result
            elif task_result.state == 'FAILURE':
                result_data['celery_error'] = str(task_result.info)
        except Exception as e:
            logger.error(f"Error getting Celery task status: {e}")

    return jsonify(result_data)


@app.route('/api/download/<int:download_id>/cancel', methods=['POST'])
@login_required
def cancel_download(download_id):
    download = Download.query.get_or_404(download_id)
    if download.user_id != current_user.id:
        return jsonify({'error': 'Unauthorized'}), 403

    if download.status in ['queued', 'downloading']:
        download.status = 'cancelled'
        db.session.commit()

        # Cancel Celery task if exists
        if download.celery_task_id:
            try:
                celery.control.revoke(download.celery_task_id, terminate=True)
                logger.info(f"Cancelled Celery task {download.celery_task_id}")
            except Exception as e:
                logger.error(f"Error cancelling Celery task: {e}")

    return jsonify({
        'success': True,
        'status': download.status
    })


@app.route('/queue')
@login_required
def queue():
    downloads = Download.query.filter_by(user_id=current_user.id).order_by(Download.created_at.desc()).all()
    return render_template('queue.html', downloads=downloads)


@app.route('/api/downloads')
@login_required
def api_downloads():
    downloads = Download.query.filter_by(user_id=current_user.id).all()
    return jsonify([{
        'id': d.id,
        'title': d.title,
        'season': d.season,
        'episode': d.episode,
        'status': d.status,
        'progress': d.progress,
        'provider': d.provider,
        'language': d.language,
        'created_at': d.created_at.isoformat(),
        'error_message': d.error_message
    } for d in downloads])


# [Alle anderen Routes bleiben gleich...]

if __name__ == '__main__':
    with app.app_context():
        db.create_all()
    app.run(debug=False, host='0.0.0.0', port=5000)  # debug=False für Production!