import os
import asyncio
import subprocess
import json
import shutil
import re
import uuid
from telethon import TelegramClient, events

# ---------- TELEGRAM API ----------
api_id = 12345678
api_hash = "your_api_hash_here"

# ---------- PATHS ----------
BASE_DIR = "/home/iman"
DOWNLOAD_DIR = f"{BASE_DIR}/telegram"
STATUS_FILE = f"{BASE_DIR}/status.json"
CONTROL_FILE = f"{BASE_DIR}/control.json"
GDRIVE_REMOTE = "gdrive:/TelegramUploads"

os.makedirs(DOWNLOAD_DIR, exist_ok=True)

client = TelegramClient("session", api_id, api_hash, timeout=60)
queue = []


# ---------- JSON HELPERS ----------
def read_json_file(path, default):
    try:
        if not os.path.exists(path):
            return default
        with open(path, "r") as f:
            return json.load(f)
    except Exception:
        return default


def write_json_file(path, data):
    tmp = f"{path}.tmp"
    with open(tmp, "w") as f:
        json.dump(data, f, indent=2)
    os.replace(tmp, path)


# ---------- INIT ----------
def init_files():
    if not os.path.exists(STATUS_FILE):
        write_json_file(STATUS_FILE, {"queue": []})

    if not os.path.exists(CONTROL_FILE):
        write_json_file(CONTROL_FILE, {
            "delete_ids": [],
            "cancel_ids": []
        })


def update_status():
    write_json_file(STATUS_FILE, {"queue": queue})


def cleanup_queue(max_items=50):
    global queue
    if len(queue) > max_items:
        queue = queue[-max_items:]


def has_enough_space(required_bytes, path=BASE_DIR):
    total, used, free = shutil.disk_usage(path)
    return free > required_bytes + 500 * 1024 * 1024


def safe_name(url: str) -> str:
    name = url.split("?")[0].rstrip("/").split("/")[-1]
    return name or "downloaded_file"


def valid_file(path: str) -> bool:
    return os.path.exists(path) and os.path.getsize(path) > 0


# ---------- CONTROL ----------
def apply_control_actions():
    global queue
    control = read_json_file(CONTROL_FILE, {"delete_ids": [], "cancel_ids": []})

    delete_ids = set(control.get("delete_ids", []))
    if delete_ids:
        queue = [item for item in queue if item.get("id") not in delete_ids]
        write_json_file(CONTROL_FILE, {
            "delete_ids": [],
            "cancel_ids": control.get("cancel_ids", [])
        })
        update_status()


def is_cancel_requested(task_id):
    control = read_json_file(CONTROL_FILE, {"delete_ids": [], "cancel_ids": []})
    return task_id in set(control.get("cancel_ids", []))


def clear_cancel_request(task_id):
    control = read_json_file(CONTROL_FILE, {"delete_ids": [], "cancel_ids": []})
    cancel_ids = [x for x in control.get("cancel_ids", []) if x != task_id]
    write_json_file(CONTROL_FILE, {
        "delete_ids": control.get("delete_ids", []),
        "cancel_ids": cancel_ids
    })


# ---------- DIRECT LINK DOWNLOAD ----------
def run_cmd(cmd):
    return subprocess.run(cmd, capture_output=True, text=True)


def download_with_wget(url, file_path, task):
    task["status"] = "downloading link (wget)"
    task["progress"] = 0
    update_status()

    result = run_cmd([
        "wget",
        "--content-disposition",
        "--trust-server-names",
        "--max-redirect=20",
        "--user-agent=Mozilla/5.0",
        "-O",
        file_path,
        url,
    ])
    if result.returncode != 0:
        raise Exception(result.stderr.strip() or result.stdout.strip() or "wget failed")


def download_with_curl(url, file_path, task):
    task["status"] = "downloading link (curl)"
    task["progress"] = 0
    update_status()

    result = run_cmd([
        "curl",
        "-L",
        "-A",
        "Mozilla/5.0",
        "-o",
        file_path,
        url,
    ])
    if result.returncode != 0:
        raise Exception(result.stderr.strip() or result.stdout.strip() or "curl failed")


def download_from_url(url, task):
    if not has_enough_space(1 * 1024 * 1024 * 1024):
        raise Exception("low disk space")

    filename = safe_name(url)
    file_path = os.path.join(DOWNLOAD_DIR, filename)
    errors = []

    # Best-effort cancel before starting
    if is_cancel_requested(task["id"]):
        clear_cancel_request(task["id"])
        raise Exception("cancelled")

    for method in (download_with_wget, download_with_curl):
        try:
            if os.path.exists(file_path):
                os.remove(file_path)

            method(url, file_path, task)

            if is_cancel_requested(task["id"]):
                clear_cancel_request(task["id"])
                if os.path.exists(file_path):
                    os.remove(file_path)
                raise Exception("cancelled")

            if valid_file(file_path):
                task["progress"] = 100
                update_status()
                return file_path

            errors.append(f"{method.__name__}: empty file")
        except Exception as e:
            errors.append(f"{method.__name__}: {e}")

    raise Exception(" | ".join(errors))


# ---------- UPLOAD ----------
def upload_with_progress(file_path, task):
    task["status"] = "uploading"
    task["progress"] = 0
    update_status()

    cmd = [
        "rclone",
        "copy",
        file_path,
        GDRIVE_REMOTE,
        "--stats", "250ms",
        "--stats-one-line",
        "--transfers", "1",
        "--checkers", "1",
        "-P",
    ]

    process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )

    last_percent = 0
    all_output = []

    for line in process.stdout:
        line = line.strip()
        all_output.append(line)
        print("RCLONE:", line, flush=True)

        if is_cancel_requested(task["id"]):
            process.kill()
            clear_cancel_request(task["id"])
            raise Exception("cancelled")

        matches = re.findall(r"(\d+)%", line)
        if matches:
            try:
                percent = max(int(x) for x in matches)
                if 0 <= percent <= 100 and percent >= last_percent:
                    last_percent = percent
                    task["progress"] = percent
                    task["status"] = "uploading"
                    update_status()
            except Exception:
                pass

    process.wait()

    if process.returncode != 0:
        last_lines = "\n".join(all_output[-10:])
        raise Exception(f"rclone upload failed with code {process.returncode}: {last_lines}")

    task["progress"] = 100
    task["status"] = "done"
    update_status()


# ---------- BACKGROUND CONTROL WATCHER ----------
async def control_watcher():
    while True:
        apply_control_actions()
        await asyncio.sleep(1)


# ---------- STARTUP ----------
init_files()


# ---------- MAIN HANDLER ----------
@client.on(events.NewMessage(chats="me", incoming=True))
async def handler(event):
    text = (event.raw_text or "").strip()
    print("Received:", repr(text), flush=True)

    # ---------- URL CASE ----------
    if text.startswith("http://") or text.startswith("https://"):
        task = {
            "id": str(uuid.uuid4()),
            "name": text,
            "size": 0,
            "status": "queued",
            "progress": 0
        }

        queue.append(task)
        cleanup_queue()
        update_status()

        try:
            file_path = await asyncio.to_thread(download_from_url, text, task)

            if is_cancel_requested(task["id"]):
                clear_cancel_request(task["id"])
                if os.path.exists(file_path):
                    os.remove(file_path)
                raise Exception("cancelled")

            task["name"] = os.path.basename(file_path)
            task["size"] = os.path.getsize(file_path)
            task["status"] = "uploading"
            task["progress"] = 0
            update_status()

            await asyncio.to_thread(upload_with_progress, file_path, task)

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

            print("URL uploaded:", file_path, flush=True)

        except Exception as e:
            msg = str(e)
            task["status"] = "cancelled" if "cancelled" in msg else f"error: {msg}"
            update_status()
            print("URL error:", e, flush=True)

        return

    # ---------- TELEGRAM FILE CASE ----------
    if not event.message.file:
        print("Ignored: not a file and not a link", flush=True)
        return

    file_name = event.message.file.name or "unknown"
    file_size = event.message.file.size or 0

    task = {
        "id": str(uuid.uuid4()),
        "name": file_name,
        "size": file_size,
        "status": "queued",
        "progress": 0
    }

    queue.append(task)
    cleanup_queue()
    update_status()

    if file_size and not has_enough_space(file_size):
        task["status"] = "error: not enough disk space"
        update_status()
        print("Telegram file error: not enough disk space", flush=True)
        return

    def progress(current, total):
        if is_cancel_requested(task["id"]):
            raise Exception("cancelled")

        if total > 0:
            task["progress"] = int(current * 100 / total)
        task["status"] = "downloading telegram file"
        update_status()

    try:
        file_path = await event.message.download_media(
            file=DOWNLOAD_DIR,
            progress_callback=progress,
        )

        if is_cancel_requested(task["id"]):
            clear_cancel_request(task["id"])
            if file_path and os.path.exists(file_path):
                os.remove(file_path)
            raise Exception("cancelled")

        task["status"] = "uploading"
        task["progress"] = 0
        update_status()

        await asyncio.to_thread(upload_with_progress, file_path, task)

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

        print("Telegram file uploaded:", file_path, flush=True)

    except Exception as e:
        msg = str(e)
        task["status"] = "cancelled" if "cancelled" in msg else f"error: {msg}"
        update_status()
        print("Telegram file error:", e, flush=True)


async def main():
    print("Listening for files and links...", flush=True)
    asyncio.create_task(control_watcher())
    await client.run_until_disconnected()


with client:
    client.loop.run_until_complete(main())