#!/usr/bin/env -S uv run --script
# /// script
# dependencies = ["psycopg[binary]", "click", "httpx"]
# ///

import argparse
import contextlib
import dataclasses
import datetime
import hashlib
import itertools
import json
import logging
import os
import re
import shlex
import shutil
import socket
import subprocess
import threading
import time
import typing
from functools import partial, wraps
from os import getenv
from pathlib import Path
from typing import Any, Callable
from urllib.parse import parse_qs
from wsgiref.simple_server import WSGIRequestHandler, make_server

import click
import httpx
import psycopg
from psycopg import sql as pgsql

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")

HandlerFunc = Callable[["Request"], "Response"]


@dataclasses.dataclass
class FfprobeResult:
    duration_sec: float
    duration_human: str
    codec: str
    fps: float
    size_bytes: int
    width: int
    height: int
    bitrate: int
    container: str
    ar: float
    sample_ar: float
    tags: dict[str, str]

    @property
    def resolution(self) -> str:
        return f"{self.width}x{self.height}"


def ffprobe(video_path: Path) -> FfprobeResult:
    proc = subprocess.run(
        args=[
            "ffprobe",
            "-v",
            "error",
            "-show_streams",
            "-show_format",
            "-print_format",
            "json",
            str(video_path),
        ],
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    try:
        proc.check_returncode()
    except subprocess.CalledProcessError as e:
        logging.error(f"failed to run ffprobe. stderr={e.stderr}")
        raise
    output: dict = json.loads(proc.stdout)
    video_stream: dict = [
        s for s in output.get("streams", []) if s.get("codec_type") == "video"
    ][0]
    video_format: dict = output.get("format", {})
    width, height = video_stream.get("width"), video_stream.get("height")
    codec = video_stream.get("codec_name")
    duration = float(video_format.get("duration", 0))
    fps_long = eval(video_stream.get("r_frame_rate", "0"))
    fps = float(f"{fps_long:.3f}") if fps_long else 0.0
    size = int(video_format.get("size", 0))
    duration_time = datetime.timedelta(seconds=int(duration))
    bitrate = int(video_stream.get("bit_rate", video_format.get("bit_rate", 0)))
    tags = video_format.get("tags", {})

    try:
        sample_w, sample_h = list(
            map(int, video_stream["sample_aspect_ratio"].split(":"))
        )
        sample_ar = (sample_w / sample_h) or 1.0
    except:
        sample_ar = 1.0

    return FfprobeResult(
        duration_sec=duration,
        duration_human=str(duration_time),
        codec=codec,
        fps=fps,
        size_bytes=size,
        width=width,
        bitrate=bitrate,
        height=height,
        container=video_path.suffix.lstrip(".").lower(),
        ar=width / height if height else 1.0,
        sample_ar=sample_ar,
        tags=tags,
    )


def hash_partial(f: Path) -> str:
    sha1 = hashlib.sha1()
    chunk_size = 1024 * 1024 * 10  # 10MB chunk size

    total_read = 0
    with f.open("rb") as file:
        while chunk := file.read(chunk_size):
            total_read += len(chunk)
            sha1.update(chunk)
            break  # Only reads the first 10MB

    return f"sha1:{total_read}:{sha1.hexdigest()}"


def chunked(iterable, n):
    it = iter(iterable)
    while True:
        chunk = tuple(itertools.islice(it, n))
        if not chunk:
            return
        yield chunk


def retry(max_attempts: int = 3, delay: float = 1.0):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_exception = None
            for attempt in range(1, max_attempts + 1):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    last_exception = e
                    if attempt < max_attempts:
                        logging.warning(
                            f"Attempt {attempt} failed: {e}. Retrying in {delay}s..."
                        )
                        time.sleep(delay)
                    else:
                        logging.error(f"All {max_attempts} attempts failed.")
            raise last_exception

        return wrapper

    return decorator


def video_exists(con: psycopg.Connection, file_path: str) -> bool:
    stmt = con.execute(
        "UPDATE videos SET last_seen_at = NOW() WHERE file_path = %s RETURNING true AS exists",
        [file_path],
    )
    row = stmt.fetchone()
    if not row:
        return False
    return row["exists"]


def save_video(
    con: psycopg.Connection,
    category: str,
    file_path: str,
    ffprobe_data: dict,
    video_hash: str,
):
    sql = """
      INSERT INTO videos(category, file_path, ffprobe, hash_partial)
      VALUES (%s, %s, %s, %s)
      ON CONFLICT(file_path) DO UPDATE SET
        last_seen_at = NOW();"""
    con.execute(
        sql,
        [
            category,
            file_path,
            json.dumps(ffprobe_data),
            video_hash,
        ],
    )


def save_error(con: psycopg.Connection, file_path: str, stderr: str):
    sql = """INSERT INTO scan_state(file_path, stderr) VALUES (%s, %s) ON conflict do nothing"""
    con.execute(
        sql,
        [
            file_path,
            stderr,
        ],
    )


def category_for_file(file_path: Path) -> str:
    s = str(file_path)
    if s.startswith("/Volumes/"):
        return f"downloaded.{file_path.parts[1].lower()}"
    if s.startswith("/mnt/box"):
        return "box"
    return "local"


def scan_videos(con: psycopg.Connection, video_paths: list[Path]) -> None:
    new_videos = []
    for i, f in enumerate(video_paths):
        if "/_picks/" in str(f):
            continue
        progress = f"[{i + 1}/{len(video_paths)}]"
        if f.is_dir():
            # Recursively find files in directory
            scan_videos(con, list(f.glob("**/*")))
            continue

        if f.suffix.lower() not in [".mp4", ".mkv", ".avi", ".mov", ".wmv", ".webm"]:
            continue

        file_path = str(f.absolute())

        if video_exists(con, file_path=file_path):
            logging.debug(f"{progress} video already exists {f=}")
            continue

        logging.info(f"{progress} probing {f=}")
        try:
            probed = ffprobe(f)
        except Exception as e:
            logging.error(f"Failed to probe {f}: {e}")
            save_error(con, file_path=file_path, stderr=str(e))
            continue

        file_hash = hash_partial(f)

        save_video(
            category=category_for_file(f),
            con=con,
            file_path=file_path,
            ffprobe_data=probed.__dict__,
            video_hash=file_hash,
        )
        logging.info(f"{progress} saved: {f.name}")
        new_videos.append(file_path)
        con.commit()

    if new_videos:
        parse_releases(con, new_videos)


def parse_releases(con: psycopg.Connection, file_paths: list[str]) -> None:
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        logging.warning("OPENROUTER_API_KEY not set, skipping filename parsing")
        return

    # Chunk file_paths by 20 items to avoid prompt token limits
    for i, chunk in enumerate(chunked(file_paths, 20), 1):
        filenames = [Path(p).name for p in chunk]
        logging.info(
            f"Parsing filenames for {len(filenames)} videos (chunk {i}) via OpenRouter..."
        )

        try:
            parsed_data = parse_filenames_with_ai(filenames, api_key)
            for item in parsed_data:
                actors = item.get("actors")
                if not isinstance(actors, list) or len(actors) == 0:
                    continue

                fname = item.get("filename")
                if not fname:
                    continue

                # Find the full path that matches this filename in the current chunk
                full_path = next((p for p in chunk if Path(p).name == fname), None)
                if not full_path:
                    continue

                # Update the release column
                release_info = {k: v for k, v in item.items() if k != "filename"}
                con.execute(
                    "UPDATE videos SET release = %s WHERE file_path = %s",
                    [json.dumps(release_info), full_path],
                )
            con.commit()
            logging.info(f"Successfully updated release info for chunk {i}")
        except Exception as e:
            logging.error(
                f"Failed to parse filenames or update database for chunk {i}: {e}"
            )


@retry(max_attempts=3, delay=0.1)
def parse_filenames_with_ai(filenames: list[str], api_key: str) -> list[dict]:
    system_prompt = """you are an expert in parsing file names.

your task is to parse the actors, studio and date, title out of filenames give me a JSON array with these fields [{filename, studio, released_at, title, actors: [...]}}]. Omit the missing / empty fields.

Actor names are usually 2 words (name and last name) but sometimes they only contain a single word. Ignore male names.
Title is the remaining part after studio, date, actors; and don't usually contain the actor names. Ignore the quality and category indicators.

output only valid JSON without any wrappers or quotes.

for example:
InTheCrack.E1890.Casey.Norhman.Provence.XXX.1080p.HEVC.x265.PRT.mp4 -> studio=InTheCrack, actors=["Casey Norhman"], title=E1890
BlackedRaw.26.05.16.Agatha.Vega.And.Ella.Hughes.Knockout.Babes.Fuck.Two.Cops.On.Duty.XXX.1080p.HEVC.x265.PRT.torrent -> {studio=BlackedRaw, actors=["Agatha Vega", "Ella Hughes"], title="Knockout Babes Fuck Two Cops On Duty", released_at=2026-05-16}
StepSiblingsCaught.26.05.14.Nata.Gold.XXX.720p.HEVC.x265.PRT.mp4	-> studio=StepSiblingsCaught, released_at=2026-05-14, actors=["Nata Gold"]
HookupHotshot.26.02.06.Episode.453.Shrooms.Q.XXX.720p.HEVC.x265.PRT.mp4 -> studio=HookupHotshot, released_at=2026-02-06, actors=["Shrooms Q"], title="Episode 453"
"""

    user_prompt = "\n".join(filenames)

    response = httpx.post(
        "https://openrouter.ai/api/v1/chat/completions",
        headers={
            "Authorization": f"Bearer {api_key}",
        },
        json={
            "model": "mistralai/ministral-3b-2512",
            "prompt_cache_key": "file_parsing",
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            "temperature": 0,
        },
        timeout=20,
    )
    response.raise_for_status()
    result = response.json()
    content = result["choices"][0]["message"]["content"].strip()

    # Remove potential markdown code blocks
    if content.startswith("```"):
        content = re.sub(r"^```(?:json)?\s*|\s*```$", "", content, flags=re.MULTILINE)

    parsed_data = json.loads(content)
    if not isinstance(parsed_data, list):
        parsed_data = [parsed_data]

    # Clean output: omit empty strings and empty arrays
    cleaned_data = []
    for entry in parsed_data:
        cleaned_entry = {
            k: v
            for k, v in entry.items()
            if v != "" and not (isinstance(v, list) and not v)
        }
        actors = cleaned_entry.get("actors", [])
        if isinstance(actors, list) and len(actors) > 0:
            actors = [a for a in actors if a and "@" not in a]
            cleaned_entry["actors"] = actors
            cleaned_data.append(cleaned_entry)

    return cleaned_data


@dataclasses.dataclass
@dataclasses.dataclass
class Request:
    path: str
    method: str
    payload: dict[str, Any]
    query: dict[str, str]
    environ: dict[str, Any]

    @classmethod
    def from_environ(cls, environ: dict) -> "Request":
        method = environ["REQUEST_METHOD"].upper()

        query = parse_qs(environ.get("QUERY_STRING", ""), keep_blank_values=True)

        payload = {}
        content_type = environ.get("CONTENT_TYPE", "")
        if method == "POST" and "application/json" in content_type:
            try:
                length = int(environ.get("CONTENT_LENGTH", 0))
                if length > 0:
                    payload = json.loads(environ["wsgi.input"].read(length))
            except (ValueError, json.JSONDecodeError):
                pass

        return cls(
            path=environ["PATH_INFO"],
            method=method,
            payload=payload,
            query={k: v[0] for k, v in query.items()},
            environ=environ,
        )


@dataclasses.dataclass
class Response:
    status: str
    headers: list[tuple[str, str]]
    body: Any

    @classmethod
    def error(cls, message: str) -> "Response":
        return cls(
            status="400 Bad Request",
            headers=[("Content-Type", "application/json")],
            body=json.dumps({"error": message}) + "\n",
        )

    @classmethod
    def from_html(cls, html: str) -> "Response":
        return cls(
            status="200 OK",
            headers=[("Content-Type", "text/html; charset=utf-8")],
            body=html,
        )

    @classmethod
    def from_json(cls, data: Any) -> "Response":
        def _default(obj):
            if isinstance(obj, (datetime.date, datetime.datetime)):
                return obj.isoformat()
            raise TypeError(f"Object of type {type(obj)} is not JSON serializable")

        return cls(
            status="200 OK",
            headers=[("Content-Type", "application/json")],
            body=json.dumps(data, default=_default) + "\n",
        )

    @classmethod
    def from_exception(cls, e: Exception) -> "Response":
        return cls(
            status="500 Internal Server Error",
            headers=[("Content-Type", "application/json")],
            body=json.dumps({"error": str(e)}) + "\n",
        )


class TinyAPI:
    def __init__(self, handlers: dict[str, HandlerFunc]):
        self.routes = self._prepare_routes(handlers)

    def _prepare_routes(self, handlers: dict[str, HandlerFunc]):
        routes = []
        for k, h in handlers.items():
            parts = k.split(maxsplit=1)
            method = parts[0].upper()
            path = parts[1].rstrip("/") or "/"
            try:
                routes.append((method, re.compile(f"^{path}$"), h))
            except re.error as e:
                raise ValueError(f"Invalid regex pattern '{path}': {e}")
        return routes

    @staticmethod
    def find_free_port() -> int:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("", 0))
            return s.getsockname()[1]

    def __call__(self, environ: dict, start_response: Callable):
        req = Request.from_environ(environ)
        path = req.path.rstrip("/") or "/"

        handler = None
        for method, pattern, h in self.routes:
            if req.method == method and pattern.match(path):
                handler = h
                break

        if not handler:
            res = Response.from_json({"error": "Not Found"})
            res.status = "404 Not Found"
        else:
            try:
                res = handler(req)
            except Exception as e:
                logging.exception("Handler crash")
                res = Response.from_exception(e)

        body = res.body if isinstance(res.body, bytes) else res.body.encode("utf-8")
        headers = res.headers + [("Content-Length", str(len(body)))]
        start_response(res.status, headers)
        return [body]

    @contextlib.contextmanager
    def serve(self, host: str = "localhost", port: int = 0):
        if port == 0:
            port = self.find_free_port()

        class LoggedRequestHandler(WSGIRequestHandler):
            def log_message(self, format: str, *args: Any) -> None:
                # args usually contains (request_line, status_code, size)
                # We redirect to our configured logger instead of sys.stderr
                logging.info("%s - %s", self.address_string(), format % args)

        server = make_server(host, port, self, handler_class=LoggedRequestHandler)
        url = f"http://{host}:{port}"

        thread = threading.Thread(target=server.serve_forever, daemon=True)
        thread.start()

        logging.info(f"Serving on {url}")
        try:
            yield url
        finally:
            logging.info("Shutting down server...")
            server.shutdown()
            server.server_close()
            thread.join(timeout=5)


@contextlib.contextmanager
def connect_db() -> typing.Generator[psycopg.Connection, typing.Any, typing.Any]:
    with psycopg.connect(
        "postgres://abdus:abdus@db.abdus.dev:5444/snot?sslmode=disable"
    ) as conn:
        conn.row_factory = psycopg.rows.dict_row
        yield conn


def handle_search(req: Request, conn: psycopg.Connection) -> Response:
    query: str = req.payload.get("query", "")
    if not query:
        return Response.error("Query is empty")

    where_sql, where_params = filter_to_where(query)

    query_stmt = pgsql.SQL(
        """
        with q as (select
            id,
            filename_from_path(file_path) as file_name,
            file_path,
            actors,
            size_bytes / 1048576 as size_mb,
            (ffprobe->>'width') || 'x' || (ffprobe->>'height') as resolution,
            duration_human,
            created_at
        from videos)
        select * from q
        where {filter}
        order by created_at desc
    """
    ).format(filter=pgsql.SQL(where_sql))
    with conn.cursor() as cursor:
        rows = cursor.execute(query_stmt, where_params).fetchall()
    return Response.from_json(
        [
            {
                **row,
                "download_url": "https://u201686:T6672ICVoWAedECH@u201686.your-storagebox.de/files"
                + row["file_path"].replace("/mnt/box/files", ""),
            }
            for row in rows
        ]
    )


def tokenize_filter_groups(filter_query: str) -> list[list[str]]:
    lexer = shlex.shlex(filter_query, posix=True, punctuation_chars="|")
    lexer.commenters = ""
    lexer.whitespace_split = True
    tokens = list(lexer)

    groups: list[list[str]] = [[]]
    for token in tokens:
        if token == "|":
            if groups[-1]:
                groups.append([])
            continue
        groups[-1].append(token)

    return [group for group in groups if group]


def token_to_condition(token: str) -> tuple[str | None, list[typing.Any]]:
    if ":" in token:
        key, value = token.split(":", 1)
        key = key.strip().lower()
        value = value.strip()
        if not value:
            return None, []

        if key == "file_name":
            return "file_name ILIKE %s", [f"%{value}%"]

        if key == "size_mb":
            match = re.match(
                r"^(>=|<=|>|<|=)?\s*(\d+(?:\.\d+)?)$", value, re.IGNORECASE
            )
            if not match:
                return None, []
            op = match.group(1) or "="
            num = float(match.group(2))
            return f"size_mb {op} %s", [num]

        if key == "actor":
            return "actors @> %s::text[]", [[value]]

        return None, []

    return "file_name ILIKE %s", [f"%{token}%"]


def filter_to_where(filter_query: str) -> tuple[str, list[typing.Any]]:
    groups = tokenize_filter_groups(filter_query=filter_query)
    if not groups:
        return "TRUE", []

    or_clauses: list[str] = []
    params: list[typing.Any] = []

    for group in groups:
        and_clauses: list[str] = []
        for token in group:
            clause, clause_params = token_to_condition(token=token)
            if not clause:
                continue
            and_clauses.append(clause)
            params.extend(clause_params)
        if and_clauses:
            or_clauses.append("(" + " AND ".join(and_clauses) + ")")

    if not or_clauses:
        return "TRUE", []

    return " OR ".join(or_clauses), params


def handle_home(req: Request, query: str = "") -> Response:
    template_file = Path(__file__).parent / "snot.html"
    html = template_file.read_text()
    if query:
        injected_json = json.dumps({"query": query})
        html = f"<script>window.ENV = {injected_json}</script>\n" + html
    return Response.from_html(html)


def handle_assets(req: Request, base_path: Path) -> Response:
    asset_path = base_path / req.path.lstrip("/")
    if not asset_path.exists() or not asset_path.is_file():
        return Response(
            status="404 Not Found",
            headers=[("Content-Type", "text/plain")],
            body="Asset not found\n",
        )

    content_type = "text/plain"
    if asset_path.suffix == ".js":
        content_type = "application/javascript"
    elif asset_path.suffix == ".css":
        content_type = "text/css"
    elif asset_path.suffix in [".html", ".htm"]:
        content_type = "text/html"
    elif asset_path.suffix == ".json":
        content_type = "application/json"

    return Response(
        status="200 OK",
        headers=[("Content-Type", content_type)],
        body=asset_path.read_text(),
    )


@contextlib.contextmanager
def run_htmlpopup(address: str):
    exe_path = shutil.which("htmlpopup")
    if not exe_path:
        raise RuntimeError(
            "htmlpopup executable not found in PATH. Please install it to use the HTML popup feature."
        )
    proc = subprocess.Popen(
        [exe_path, "--title", "s·m·u·t", address], text=True, stdout=subprocess.PIPE
    )
    try:
        yield
        proc.wait()
    finally:
        proc.terminate()


@click.group(invoke_without_command=True)
@click.pass_context
def cli(ctx: click.Context):
    """Snot - Video Browser and Scanner"""
    if ctx.invoked_subcommand is None:
        ctx.invoke(serve)


@cli.command()
@click.option(
    "--port",
    type=int,
    default=int(getenv("PORT", "0")),
    help="Port to run the server on",
)
@click.option("--query", type=str, default="", help="Initial query string")
def serve(port: int, query: str):
    """Start the web UI"""
    with connect_db() as conn:
        handlers = {
            "POST /search": partial(handle_search, conn=conn),
            "GET /": partial(handle_home, query=query),
            "GET /.*": partial(handle_assets, base_path=Path(__file__).parent),
        }
        api = TinyAPI(handlers=handlers)

        with api.serve(port=port) as server_url:
            logging.info(f"Server running at {server_url}")
            logging.info("Available endpoints:")
            for path in handlers.keys():
                logging.info(f"  {path}")
            with run_htmlpopup(server_url):
                logging.info("HTML popup started. Press Ctrl+C to stop.")


@cli.command()
@click.argument("paths", nargs=-1, type=Path)
def scan(paths: list[Path]):
    """Scan video files and update the database"""
    if not paths:
        click.echo("No paths provided to scan.")
        return

    with connect_db() as con:
        scan_videos(con=con, video_paths=list(paths))


@cli.command("mark-deletion")
@click.argument("paths", nargs=-1, type=Path)
def mark_deletion(paths: list[Path]):
    """Mark video files for deletion by setting marked_for_deletion_at"""
    if not paths:
        click.echo("No paths provided.")
        return

    with connect_db() as con:
        for path in paths:
            file_name = path.name
            rows = con.execute(
                "UPDATE videos SET marked_for_deletion_at = NOW() WHERE file_name = %s RETURNING file_path",
                [file_name],
            ).fetchall()
            if rows:
                for row in rows:
                    click.echo(f"Marked for deletion: {row['file_path']}")
            else:
                click.echo(f"Not found in database: {file_name}", err=True)
                continue

            if not click.get_text_stream("stdin").isatty():
                continue

            stem = path.stem
            siblings = con.execute(
                "SELECT file_path, file_name FROM videos WHERE file_name != %s AND file_name LIKE %s AND marked_for_deletion_at IS NULL",
                [file_name, f"{stem}.%"],
            ).fetchall()
            for sibling in siblings:
                if click.confirm(f"  Also mark {sibling['file_name']}?", default=False):
                    con.execute(
                        "UPDATE videos SET marked_for_deletion_at = NOW() WHERE file_path = %s",
                        [sibling["file_path"]],
                    )
                    click.echo(f"  Marked for deletion: {sibling['file_path']}")

        con.commit()


if __name__ == "__main__":
    cli()