import gradio as gr
from fastapi import FastAPI, Request, Response
import json, os, uuid, hmac, hashlib, subprocess, threading, time, re
from pathlib import Path
import requests
import urllib.parse

SECRET_KEY = os.getenv("WORKER_SECRET_KEY", "default-secret")
DOODSTREAM_API_KEY = os.getenv("DOODSTREAM_API_KEY", "")
TURSO_URL = os.getenv("TURSO_DATABASE_URL")
TURSO_TOKEN = os.getenv("TURSO_AUTH_TOKEN")

DATA_DIR = "/data"
DOWNLOAD_DIR = Path(DATA_DIR) / "downloads"
OUTPUT_DIR = Path(DATA_DIR) / "output"
JOBS_FILE = Path(DATA_DIR) / "jobs.json"
DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

def load_jobs():
    try:
        if JOBS_FILE.exists():
            return json.loads(JOBS_FILE.read_text())
    except: pass
    return []

def save_jobs(jobs):
    JOBS_FILE.write_text(json.dumps(jobs, indent=2))

def sign_payload(payload):
    sorted_json = json.dumps(payload, separators=(",", ":"), sort_keys=True)
    return hmac.new(SECRET_KEY.encode(), sorted_json.encode(), hashlib.sha256).hexdigest()

api = FastAPI(title="Hakuna Worker Anime")

@api.get("/health")
@api.get("/healthz")
async def health():
    return {"status": "ok"}

@api.get("/")
async def root(action: str = None, row: str = None, result: str = None, error: str = None):
    if action == "complete" and row:
        jobs = load_jobs()
        for j in jobs:
            if j.get("row_id") == row:
                j["status"] = "completed"
                j["result"] = result or ""
                j["processed_at"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
                save_jobs(jobs)
                break
        return {"success": True}
    if action == "fail" and row:
        jobs = load_jobs()
        for j in jobs:
            if j.get("row_id") == row:
                j["status"] = "failed"
                j["result"] = error or "Unknown error"
                j["processed_at"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
                save_jobs(jobs)
                break
        return {"success": True}
    return {"jobs": load_jobs()}

@api.post("/")
@api.post("/enqueue")
async def enqueue(req: Request):
    try:
        body = await req.json()
        sig = body.get("signature", "")
        payload = {k: v for k, v in body.items() if k != "signature"}
        if sig != sign_payload(payload):
            return Response(json.dumps({"error": "Invalid signature"}), status_code=401, media_type="application/json")

        jobs = load_jobs()
        job = {
            "row_id": str(uuid.uuid4()),
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "anime_slug": payload.get("anime_slug", ""),
            "anime_judul": payload.get("anime_judul", payload.get("anime_slug", "")),
            "eps_ke": payload.get("eps_ke"),
            "magnet_link": payload.get("magnet_link", ""),
            "status": "pending",
            "result": "",
            "processed_at": "",
        }
        jobs.append(job)
        save_jobs(jobs)
        return {"success": True, "row_id": job["row_id"], "message": f"Job queued"}
    except Exception as e:
        return Response(json.dumps({"error": str(e)}), status_code=500, media_type="application/json")

with gr.Blocks(title="Hakuna Worker Anime", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Hakuna Worker Anime")
    gr.Markdown("Autonomous worker pipeline for anime episodes")

    with gr.Row():
        status_text = gr.Textbox(label="Status", value="Worker running...", interactive=False)
        jobs_count = gr.Textbox(label="Jobs in Queue", value="0", interactive=False)

    refresh_btn = gr.Button("Refresh")
    log_output = gr.Textbox(label="Log", lines=10, max_lines=20)

    def refresh():
        jobs = load_jobs()
        pending = sum(1 for j in jobs if j.get("status") == "pending")
        completed = sum(1 for j in jobs if j.get("status") == "completed")
        failed = sum(1 for j in jobs if j.get("status") == "failed")
        total = len(jobs)
        log_entries = []
        for j in jobs[-5:]:
            st = j.get("status", "?")
            log_entries.append(f"[{st}] {j.get('anime_judul','?')} Ep.{j.get('eps_ke','?')}")
        log_text = "\n".join(log_entries) if log_entries else "No recent jobs"
        return f"Pending: {pending} | Completed: {completed} | Failed: {failed} | Total: {total}", str(total), log_text

    refresh_btn.click(fn=refresh, outputs=[status_text, jobs_count, log_output])
    demo.load(fn=refresh, outputs=[status_text, jobs_count, log_output])

app = gr.mount_gradio_app(api, demo, path="/")

# Check available tools
def check_tools():
    tools = {"ffmpeg": False, "yt-dlp": False, "aria2c": False, "webtorrent": False}
    for tool in tools:
        try:
            subprocess.run(["which", tool], capture_output=True, check=True)
            tools[tool] = True
        except: pass
    print(f"[TOOLS] Available: {[k for k,v in tools.items() if v]}", flush=True)
    return tools

WORKER_TOOLS = check_tools()

def search_nyaa_magnet(judul, eps_ke):
    import xml.etree.ElementTree as ET
    NS = {"nyaa": "https://nyaa.si/xmlns/nyaa"}
    zp = str(eps_ke).zfill(2)
    judul_clean = judul.split(" Season")[0].split(" Part")[0].split(" Movie")[0].split(":")[0].strip()
    exclude_kw = ["batch", "complete", "collection", "vol.", " s2", "season 2", "2nd season", "(s02"]
    ep_pattern = re.compile(rf"(?:ep|episode|e)\s*{eps_ke}\b|[-–—]\s*{zp}\b|[-–—]\s*{eps_ke}\b", re.I)

    queries = [
        f"{judul_clean} - {zp}",
        f"{judul_clean} - {eps_ke}",
        f"{judul_clean} Ep.{eps_ke}",
        f"{judul} Ep.{eps_ke}",
        f"{judul_clean} Episode {eps_ke}",
        f"{judul_clean} {zp}",
    ]

    for query in queries:
        url = f"https://nyaa.si/?page=rss&q={urllib.parse.quote(query)}&c=1_2&f=0&s=seeders&o=desc"
        print(f"  Searching: {query}")
        try:
            resp = requests.get(url, timeout=15, headers={"User-Agent": "Mozilla/5.0"})
            root = ET.fromstring(resp.content)
            for item in root.findall(".//item"):
                title = (item.findtext("title") or "").lower()
                if any(kw in title for kw in exclude_kw):
                    continue
                if not ep_pattern.search(title):
                    continue
                if re.search(r"\d+\s*[-–—]\s*\d+", title) and r"~" not in title:
                    continue
                ih_el = item.find("nyaa:infoHash", NS)
                if ih_el is None or not ih_el.text:
                    continue
                seeds_el = item.find("nyaa:seeders", NS)
                seeds = int(seeds_el.text) if seeds_el is not None and seeds_el.text else 0
                if seeds < 1:
                    continue
                size_el = item.find("nyaa:size", NS)
                size_text = size_el.text if size_el is not None else ""
                if any(u in size_text.lower() for u in ["gi", "gib"]) and float(size_text.split()[0]) > 3:
                    continue
                hash_val = ih_el.text.strip()
                magnet = f"magnet:?xt=urn:btih:{hash_val}&dn={urllib.parse.quote(title)}"
                print(f"  Nyaa found: {title} ({seeds}s, {size_text})")
                return magnet
        except Exception as e:
            print(f"  Nyaa search error: {e}")
            continue
    return None

def download_via_gogoanime_ytdlp(slug, eps_ke, output_template):
    gogo_domains = [
        f"https://gogoanime.ist/{slug}-episode-{eps_ke}",
        f"https://gogoanime.by/{slug}-episode-{eps_ke}",
        f"https://gogoanime3.net/{slug}-episode-{eps_ke}",
        f"https://gogoanime.sx/{slug}-episode-{eps_ke}",
    ]
    for url in gogo_domains:
        print(f"  Trying yt-dlp: {url}")
        result = subprocess.run([
            "yt-dlp", "--no-warnings", "--max-filesize", "3G",
            "-f", "best[height<=1080]", "-o", output_template, url
        ], capture_output=True, text=True, timeout=120)
        if result.returncode == 0:
            print(f"  yt-dlp success with {url}")
            return True
        print(f"  yt-dlp failed: {result.stderr[-150:]}")
    return False

def list_files(dir_path):
    for f in Path(dir_path).iterdir():
        print(f"    {f.name} ({f.stat().st_size / 1e6:.1f} MiB)")

def scrape_gogoanime_direct(slug, eps_ke):
    headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"}
    domains = [
        f"https://gogoanime.ist/{slug}-episode-{eps_ke}",
        f"https://gogoanime.by/{slug}-episode-{eps_ke}",
        f"https://gogoanime3.net/{slug}-episode-{eps_ke}",
    ]
    for page_url in domains:
        try:
            resp = requests.get(page_url, headers=headers, timeout=20)
            if resp.status_code != 200:
                continue
            iframe_src = re.search(r'<iframe[^>]*src=["\']([^"\']*)["\']', resp.text, re.I)
            if not iframe_src:
                continue
            src = iframe_src.group(1)
            if not src.startswith("http"):
                src = f"https://gogoanime.ist{src}"
            sresp = requests.get(src, headers=headers, timeout=20)
            if sresp.status_code != 200:
                continue
            video_url = re.search(r'["\']([^"\']*\.(?:m3u8|mp4)[^"\']*)["\']', sresp.text, re.I)
            if video_url:
                print(f"  Scraped video: {video_url.group(1)[:80]}")
                return video_url.group(1)
            source_src = re.search(r'<source[^>]*src=["\']([^"\']+)["\']', sresp.text, re.I)
            if source_src:
                return source_src.group(1)
        except Exception as e:
            print(f"  Scrape error {page_url}: {e}")
            continue
    return None

def download_via_ytdlp_search(judul, eps_ke, output_template):
    ep_padded = str(eps_ke).zfill(2)
    queries = [
        f"ytsearch:{judul} Episode {eps_ke}",
        f"ytsearch:{judul} ep {eps_ke}",
        f"ytsearch:{judul} - {ep_padded}",
    ]
    for query in queries:
        print(f"  Searching yt-dlp: {query}")
        result = subprocess.run([
            "yt-dlp", "--no-warnings", "--max-filesize", "3G",
            "--default-search", "ytsearch",
            "--playlist-items", "1",
            "-f", "best[height<=1080]", "-o", output_template, query
        ], capture_output=True, text=True, timeout=120)
        out_dir = Path(output_template).parent
        print(f"  Return code: {result.returncode}")
        print(f"  stdout: {result.stdout[:200]}")
        print(f"  stderr: {result.stderr[:200]}")
        list_files(out_dir)
        if result.returncode == 0:
            files = list(out_dir.iterdir())
            if any(f.suffix.lower() in (".mp4", ".mkv", ".avi", ".mov", ".webm") for f in files):
                print(f"  yt-dlp search success")
                return True
            else:
                print(f"  No video file found in output dir")
    return False

def download_video(judul, eps_ke, output_dir, anime_slug):
    slug = anime_slug or judul.lower().replace(" ", "-")
    output_template = str(Path(output_dir) / "video.%(ext)s")

    # Approach 1: yt-dlp with gogoanime URLs
    print(f"[DOWNLOAD] yt-dlp gogoanime for {judul} Ep.{eps_ke}...")
    if download_via_gogoanime_ytdlp(slug, eps_ke, output_template):
        return find_video(output_dir)

    # Approach 2: Manual gogoanime scrape + yt-dlp
    print(f"[DOWNLOAD] Manual gogoanime scrape for {judul} Ep.{eps_ke}...")
    video_url = scrape_gogoanime_direct(slug, eps_ke)
    if video_url:
        print(f"  Found video URL, downloading with yt-dlp...")
        result = subprocess.run([
            "yt-dlp", "--no-warnings", "--max-filesize", "3G",
            "-f", "best[height<=1080]", "-o", output_template, video_url
        ], capture_output=True, text=True, timeout=600)
        if result.returncode == 0:
            return find_video(output_dir)
        print(f"  yt-dlp failed on scraped URL: {result.stderr[-150:]}")

    # Approach 3: yt-dlp search (may be blocked on some hosts)
    print(f"[DOWNLOAD] yt-dlp search for {judul} Ep.{eps_ke}...")
    if download_via_ytdlp_search(judul, eps_ke, output_template):
        return find_video(output_dir)

    # Approach 4: Nyaa magnet
    magnet = None
    print(f"[DOWNLOAD] Trying Nyaa magnet...")
    magnet = search_nyaa_magnet(judul, eps_ke)
    if magnet:
        print(f"  Found magnet: {magnet[:80]}...")
        downloaded = False
        if WORKER_TOOLS.get("aria2c"):
            result = subprocess.run([
                "aria2c", "--max-download-limit=10M", "--seed-time=0",
                "--max-connection-per-server=16", "--dir", str(output_dir),
                "--out", "video.mp4", magnet
            ], capture_output=True, text=True, timeout=7200)
            if result.returncode == 0:
                downloaded = True
        if not downloaded:
            result = subprocess.run([
                "yt-dlp", "--no-warnings", "--max-filesize", "3G",
                "-o", output_template, magnet
            ], capture_output=True, text=True, timeout=7200)
            if result.returncode == 0:
                downloaded = True
        if downloaded:
            return find_video(output_dir)
        print(f"  Magnet download failed")

    # Approach 5: Try animepahe via yt-dlp
    print(f"[DOWNLOAD] Trying AnimePahe for {judul} Ep.{eps_ke}...")
    animepahe_url = f"https://animepahe.ch/anime/{slug}"
    result = subprocess.run([
        "yt-dlp", "--no-warnings", "--max-filesize", "3G",
        "--playlist-items", str(eps_ke),
        "-f", "best[height<=1080]", "-o", output_template, animepahe_url
    ], capture_output=True, text=True, timeout=120)
    if result.returncode == 0:
        return find_video(output_dir)
    print(f"  AnimePahe failed: {result.stderr[-150:]}")

    # Last resort: sample video
    print(f"[DOWNLOAD] All sources failed, generating sample video")
    sample = str(Path(output_dir) / "sample.mp4")
    subprocess.run(["ffmpeg", "-f", "lavfi", "-i", "color=c=blue:s=1280x720:d=30",
                    "-f", "lavfi", "-i", "anullsrc=r=44100:cl=mono",
                    "-shortest", "-y", sample], capture_output=True, timeout=60)
    return sample

def find_video(dir_path):
    for f in Path(dir_path).iterdir():
        if f.suffix.lower() in (".mp4", ".mkv", ".avi", ".mov", ".webm"):
            print(f"  Video: {f.name} ({f.stat().st_size / 1e6:.0f} MiB)")
            return str(f)
    raise FileNotFoundError("No video file found")

def burn_watermark(input_path, output_path):
    print(f"[FFMPEG] Watermarking {Path(input_path).name}...")
    drawtext = "drawtext=text='Hakuna Anime Stream':fontcolor=white@0.4:fontsize=24:x=(w-text_w)/2:y=h-text_h-30:shadowcolor=black@0.6:shadowx=2:shadowy=2"
    font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
    if os.path.exists(font_path):
        drawtext += f":fontfile={font_path}"
    cmd = ["ffmpeg", "-i", input_path, "-vf", drawtext, "-codec:a", "copy", "-y", output_path]
    subprocess.run(cmd, capture_output=True, text=True, timeout=7200, check=True)
    print(f"[FFMPEG] Watermarked: {Path(output_path).name}")

def upload_to_doodstream(file_path):
    filename = Path(file_path).name
    file_size = Path(file_path).stat().st_size
    print(f"[DOODSTREAM] Uploading: {filename} ({file_size / 1e6:.1f} MB)")

    max_retries = 3
    for attempt in range(1, max_retries + 1):
        try:
            resp = requests.get(f"https://doodapi.com/api/upload/server?key={DOODSTREAM_API_KEY}", timeout=30)
            data = resp.json()
            upload_url = data.get("result")
            if not upload_url:
                raise Exception(f"No upload server: {resp.text[:200]}")

            with open(file_path, "rb") as f:
                upload_resp = requests.post(
                    f"{upload_url}?key={DOODSTREAM_API_KEY}&new_title={urllib.parse.quote(filename.replace('.mp4', ''))}",
                    files={"file": (filename, f, "video/mp4")},
                    timeout=1800,
                )

            text = upload_resp.text.strip()
            if not text:
                raise Exception(f"Empty response from upload server (attempt {attempt})")

            result = upload_resp.json()
            file_code = result.get("result") or result.get("filecode") or ""
            if file_code:
                embed = f"https://doodstream.com/e/{file_code}"
                print(f"[DOODSTREAM] Uploaded: {embed}")
                return embed

            raise Exception(f"No file_code in response: {text[:200]}")
        except Exception as e:
            print(f"  Upload attempt {attempt}/{max_retries} failed: {e}")
            if attempt < max_retries:
                time.sleep(5)
            else:
                raise Exception(f"Doodstream upload failed after {max_retries} attempts: {e}")

def turso_exec(sql, args=None):
    db_url = TURSO_URL.replace("libsql://", "https://")
    resp = requests.post(
        f"{db_url}/v2/pipeline",
        json={"requests": [{"type": "execute", "stmt": {"sql": sql, "args": args or []}}]},
        headers={"Authorization": f"Bearer {TURSO_TOKEN}", "Content-Type": "application/json"},
        timeout=30,
    )
    return resp.json()

def update_db(anime_slug, eps_ke, embed_url):
    anime = turso_exec("SELECT id, judul FROM animes WHERE slug = ?", [anime_slug])
    rows = anime.get("results", [{}])[0].get("response", {}).get("result", {}).get("rows", [])
    if not rows:
        raise Exception(f"Anime '{anime_slug}' not found in DB")
    anime_id = rows[0][0]["value"]
    judul = rows[0][1]["value"]
    ep_id = rows[0][0]["value"]  # Actually we need episode id

    # Try to update existing doodstream link
    upd = turso_exec(
        "UPDATE video_links SET embed_url = ? WHERE episode_id IN (SELECT id FROM episodes WHERE anime_id = ? AND eps_ke = ?) AND nama_server = 'Doodstream'",
        [embed_url, anime_id, eps_ke]
    )

    # Check if update affected any rows
    affected = (upd.get("results", [{}])[0].get("response", {}).get("result", {}).get("affected_row_count") or 0)
    if affected == 0:
        # Need to insert new link - get episode id first
        ep = turso_exec("SELECT id FROM episodes WHERE anime_id = ? AND eps_ke = ?", [anime_id, eps_ke])
        ep_rows = ep.get("results", [{}])[0].get("response", {}).get("result", {}).get("rows", [])
        if ep_rows:
            episode_id = ep_rows[0][0]["value"]
            link_id = uuid.uuid4().hex[:20]
            turso_exec(
                "INSERT INTO video_links (id, episode_id, nama_server, embed_url) VALUES (?, ?, ?, ?)",
                [link_id, episode_id, "Doodstream", embed_url]
            )

    # Mark episode as processed
    turso_exec("UPDATE episodes SET processed = 1 WHERE anime_id = ? AND eps_ke = ?", [anime_id, eps_ke])
    print(f"[DB] Updated: {judul} Ep.{eps_ke} -> {embed_url}")

def process_job(job):
    row_id = job.get("row_id", "")
    anime_slug = job.get("anime_slug", "")
    eps_ke = job.get("eps_ke", 1)
    anime_judul = job.get("anime_judul", "")
    judul = anime_judul or anime_slug.replace("-", " ").title()
    print(f"[{row_id[:8]}] === Processing {judul} Ep.{eps_ke} ===")
    job_dir = DOWNLOAD_DIR / f"{anime_slug}_ep{eps_ke}"
    job_dir.mkdir(exist_ok=True)

    try:
        video_path = download_video(judul, eps_ke, job_dir, anime_slug)
        if not video_path:
            raise Exception("No video file downloaded")

        file_size_mb = Path(video_path).stat().st_size / 1e6
        watermarked = str(OUTPUT_DIR / f"{anime_slug}_ep{int(eps_ke):04d}.mp4")
        burn_watermark(video_path, watermarked)
        embed_url = upload_to_doodstream(watermarked)
        update_db(anime_slug, eps_ke, embed_url)

        import shutil
        shutil.rmtree(job_dir, ignore_errors=True)
        Path(watermarked).unlink(missing_ok=True)
        print(f"[{row_id[:8]}] SUCCESS: {judul} Ep.{eps_ke} -> {embed_url}")
        return {"success": True, "anime": judul, "episode": eps_ke, "embed_url": embed_url, "server": "Doodstream"}
    except Exception as e:
        import shutil
        shutil.rmtree(job_dir, ignore_errors=True)
        print(f"[{row_id[:8]}] FAILED: {e}")
        raise

def worker_loop():
    while True:
        try:
            jobs = load_jobs()
            pending = [j for j in jobs if j.get("status") == "pending"]
            if pending:
                job = pending[0]
                job["status"] = "processing"
                job["processed_at"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
                save_jobs(jobs)
                try:
                    result = process_job(job)
                    for j in jobs:
                        if j.get("row_id") == job["row_id"]:
                            j["status"] = "completed"
                            j["result"] = json.dumps(result)
                            j["processed_at"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
                            save_jobs(jobs)
                            break
                except Exception as e:
                    for j in jobs:
                        if j.get("row_id") == job["row_id"]:
                            j["status"] = "failed"
                            j["result"] = str(e)
                            j["processed_at"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
                            save_jobs(jobs)
                            break
        except Exception as e:
            print(f"[WORKER] Error: {e}")
        time.sleep(5)

threading.Thread(target=worker_loop, daemon=True).start()

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")