import base64
import ipaddress
import json
import logging
import os
import socket
import threading
from contextlib import contextmanager
from io import BytesIO
from urllib.parse import urlsplit
from uuid import uuid4

from PIL import Image, ImageFilter, ImageOps, ImageStat

from pigway_frame_flow.core.photo_preferences import flags
from pigway_frame_flow.errors import ErrorCode as E
from pigway_frame_flow.errors import FlowError
from pigway_frame_flow.storage.database import now_iso

MAX_IMAGE_PIXELS = 25_000_000
Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS


def convert(
    data,
    fit="contain",
    profile="color",
    *,
    size=(1200, 1600),
    centering=(0.5, 0.5),
    quality=93,
):
    try:
        with Image.open(BytesIO(data), formats=("JPEG", "PNG", "WEBP")) as source:
            if source.format not in {"JPEG", "PNG", "WEBP"}:
                raise ValueError()
            if source.width * source.height > MAX_IMAGE_PIXELS:
                raise ValueError()
            img = ImageOps.exif_transpose(source).convert("RGBA")
            canvas = Image.new("RGBA", img.size, "white")
            canvas.alpha_composite(img)
            img = canvas.convert("RGB")
            if fit == "contain":
                result = ImageOps.pad(
                    img, size, method=Image.Resampling.LANCZOS, color="white"
                )
            else:
                center = centering
                if fit == "smart":
                    # Lightweight edge saliency: slide an aspect-correct window
                    # over a thumbnail; retain the region with most detail.
                    thumb = img.copy()
                    thumb.thumbnail((240, 240))
                    edges = thumb.convert("L").filter(ImageFilter.FIND_EDGES)
                    w, h = edges.size
                    cw, ch = (
                        (int(h * size[0] / size[1]), h)
                        if w / h > size[0] / size[1]
                        else (w, int(w * size[1] / size[0]))
                    )
                    choices = []
                    for step in range(11):
                        x, y = int((w - cw) * step / 10), int((h - ch) * step / 10)
                        score = ImageStat.Stat(edges.crop((x, y, x + cw, y + ch))).sum[
                            0
                        ]
                        # Favor centered crops in ties, suppress border artifacts.
                        choices.append((score - abs(step - 5) * 0.01, x, y))
                    _, x, y = max(choices)
                    center = (
                        (x / (w - cw)) if w != cw else 0.5,
                        (y / (h - ch)) if h != ch else 0.5,
                    )
                result = ImageOps.fit(
                    img, size, method=Image.Resampling.LANCZOS, centering=center
                )
            if profile == "grayscale":
                result = ImageOps.grayscale(result).convert("RGB")
            out = BytesIO()
            result.save(out, "JPEG", quality=quality, optimize=True)
            return out.getvalue()
    except (OSError, ValueError, Image.DecompressionBombError):
        raise FlowError(E.INVALID_REQUEST, "Unsupported or invalid image") from None


class Assets:
    def __init__(self, db, settings):
        self.db, self.settings = db, settings
        self.file_lock = threading.RLock()
        db.execute(
            "CREATE TABLE IF NOT EXISTS asset_file_journal(id TEXT PRIMARY KEY,kind TEXT NOT NULL)"
        )
        db.execute(
            "CREATE TABLE IF NOT EXISTS asset_file_moves(id TEXT PRIMARY KEY,source_kind TEXT NOT NULL,target_kind TEXT NOT NULL)"
        )
        self.recover_files()

    @staticmethod
    def sync_directory(directory):
        fd = os.open(directory, os.O_RDONLY)
        try:
            os.fsync(fd)
        finally:
            os.close(fd)

    def recover_files(self):
        with self.file_lock:
            self._recover_files()

    def _recover_files(self):
        for row in self.db.all("SELECT * FROM asset_file_journal"):
            path = self.settings.data_dir / row["kind"] / (row["id"] + ".jpg")
            try:
                if not self.db.one("SELECT id FROM assets WHERE id=?", (row["id"],)):
                    path.unlink(missing_ok=True)
                path.with_suffix(".part").unlink(missing_ok=True)
                if path.parent.exists():
                    self.sync_directory(path.parent)
                self.db.execute(
                    "DELETE FROM asset_file_journal WHERE id=?", (row["id"],)
                )
            except OSError:
                logging.getLogger(__name__).warning("Asset file cleanup pending")

        for row in self.db.all("SELECT * FROM asset_file_moves"):
            asset = self.db.one("SELECT kind FROM assets WHERE id=?", (row["id"],))
            obsolete = (
                row["source_kind"]
                if asset and asset["kind"] == row["target_kind"]
                else row["target_kind"]
            )
            path = self.settings.data_dir / obsolete / (row["id"] + ".jpg")
            try:
                path.unlink(missing_ok=True)
                self.sync_directory(path.parent)
                self.db.execute("DELETE FROM asset_file_moves WHERE id=?", (row["id"],))
            except OSError:
                logging.getLogger(__name__).warning("Asset promotion cleanup pending")

    @contextmanager
    def creating(self, data, name="image.jpg", kind="images"):
        with self.file_lock:
            with self._creating(data, name, kind) as asset:
                yield asset

    @contextmanager
    def _creating(self, data, name="image.jpg", kind="images"):
        from pathlib import Path

        if kind not in {"images", "generated", "cache", "source-images", "renditions"}:
            raise ValueError("Invalid asset kind")
        name = Path(name).stem[:110] + ".jpg"
        asset_id = uuid4().hex
        path = self.settings.data_dir / kind / (asset_id + ".jpg")
        path.parent.mkdir(exist_ok=True)
        self.db.execute("INSERT INTO asset_file_journal VALUES(?,?)", (asset_id, kind))
        temporary = path.with_suffix(".part")
        try:
            with temporary.open("xb") as handle:
                handle.write(data)
                handle.flush()
                os.fsync(handle.fileno())
            temporary.replace(path)
            self.sync_directory(path.parent)
            with self.db.transaction():
                self.db.execute(
                    "INSERT INTO assets VALUES(?,?,?,?)",
                    (asset_id, name, kind, now_iso()),
                )
                yield self.info(asset_id)
                self.db.execute(
                    "DELETE FROM asset_file_journal WHERE id=?", (asset_id,)
                )
        except BaseException:
            # The journal remains durable if filesystem cleanup also fails.
            self.recover_files()
            raise

    def put(self, data, name="image.jpg", kind="images"):
        with self.creating(data, name, kind) as asset:
            return asset

    def record_origin(self, asset_id, operation):
        job = self.db.one(
            "SELECT jobs.id,jobs.payload FROM job_runs JOIN jobs ON jobs.id=job_runs.job_id WHERE operation_id=?",
            (operation["id"],),
        )
        import json

        payload = json.loads(operation["payload"])
        content = payload.get("content", payload)
        origin = (
            "task"
            if job
            else "template"
            if content.get("type") in ("template", "auto_collage")
            else "manual"
            if operation["client_id"].startswith("user:")
            else "external"
        )
        self.db.execute(
            "INSERT OR IGNORE INTO asset_metadata(asset_id,origin,job_id,job_name) VALUES(?,?,?,?)",
            (
                asset_id,
                origin,
                job["id"] if job else None,
                json.loads(job["payload"])["name"] if job else None,
            ),
        )

    def promote(self, asset_id):
        with self.file_lock:
            row = self.db.one("SELECT * FROM assets WHERE id=?", (asset_id,))
            if row and row["kind"] == "cache":
                source = self.db.one(
                    "SELECT payload FROM asset_sources WHERE asset_id=?", (asset_id,)
                )
                kind = (
                    "generated"
                    if source
                    and json.loads(source["payload"]).get("type") == "template"
                    else "images"
                )
                target = self.settings.data_dir / kind / (asset_id + ".jpg")
                self.recover_files()
                self.db.execute(
                    "INSERT INTO asset_file_moves VALUES(?,?,?)",
                    (asset_id, "cache", kind),
                )
                try:
                    # Keep the source valid until the database points to the durable target.
                    os.link(self.path(asset_id), target)
                    self.sync_directory(target.parent)
                    self.db.execute(
                        "UPDATE assets SET kind=? WHERE id=?", (kind, asset_id)
                    )
                except BaseException:
                    self.recover_files()
                    raise
                self.recover_files()
            return self.info(asset_id)

    def info(self, asset_id):
        row = self.db.one("SELECT * FROM assets WHERE id=?", (asset_id,))
        if not row:
            raise FlowError(E.NOT_FOUND, "Asset not found", 404)
        meta = self.db.one("SELECT * FROM asset_metadata WHERE asset_id=?", (asset_id,))
        provenance = self.db.one(
            "SELECT payload FROM asset_sources WHERE asset_id=?", (asset_id,)
        )
        return {
            "source_path": json.loads(provenance["payload"]).get("source_path")
            if provenance
            else None,
            "starred": bool(meta and meta["starred"]),
            "excluded": flags(self.db, "@assets", asset_id)["excluded"],
            "origin": meta["origin"] if meta else "unknown",
            "job_id": meta["job_id"] if meta else None,
            "job_name": meta["job_name"] if meta else None,
            "asset_id": row["id"],
            "source": "generated"
            if row["kind"] in ("generated", "renditions")
            else "uploaded",
            "name": row["name"],
            "url": f"/api/v1/assets/{row['id']}",
            "created_at": row["created_at"],
        }

    def path(self, asset_id):
        row = self.db.one("SELECT * FROM assets WHERE id=?", (asset_id,))
        if not row:
            raise FlowError(E.NOT_FOUND, "Asset not found", 404)
        return self.settings.data_dir / row["kind"] / (row["id"] + ".jpg")

    def decode(self, value):
        try:
            data = base64.b64decode(value, validate=True)
        except ValueError:
            raise FlowError(E.INVALID_REQUEST, "Invalid base64 image") from None
        if len(data) > self.settings.max_bytes:
            raise FlowError(E.INVALID_REQUEST, "Image exceeds 20 MiB")
        return data

    def fetch_url(
        self,
        url,
        *,
        method="GET",
        headers=None,
        body=None,
        limit=None,
        accepted_status=(200,),
        timeout=10,
    ):
        limit = limit or self.settings.max_bytes
        parsed = urlsplit(url)
        if (
            parsed.scheme != "https"
            or not parsed.hostname
            or parsed.username
            or parsed.password
            or parsed.port not in (None, 443)
        ):
            raise FlowError(E.INVALID_REQUEST, "Image URL requires HTTPS on port 443")
        if parsed.hostname.lower() not in self.settings.image_url_hosts:
            raise FlowError(E.FORBIDDEN, "Image host is not allowed", 403)
        try:
            addresses = socket.getaddrinfo(
                parsed.hostname, 443, type=socket.SOCK_STREAM
            )
            if not addresses or any(
                not ipaddress.ip_address(x[4][0]).is_global for x in addresses
            ):
                raise FlowError(
                    E.FORBIDDEN, "Private image addresses are forbidden", 403
                )
            # Pin the validated resolution for this request, including TLS SNI,
            # instead of performing a second DNS lookup (rebinding).
            import http.client
            import ssl

            ip = str(addresses[0][4][0])

            class PinnedHTTPS(http.client.HTTPSConnection):
                def connect(self):
                    sock = socket.create_connection((ip, 443), timeout=timeout)
                    self.sock = ssl.create_default_context().wrap_socket(
                        sock, server_hostname=parsed.hostname
                    )

            connection = PinnedHTTPS(parsed.hostname, timeout=timeout)
            try:
                connection.request(
                    method,
                    (parsed.path or "/") + ("?" + parsed.query if parsed.query else ""),
                    headers=headers or {},
                    body=body,
                )
                response = connection.getresponse()
                if response.status not in accepted_status:
                    raise FlowError(
                        E.INVALID_REQUEST, "Image URL returned an invalid response"
                    )
                data = response.read(limit + 1)
                if len(data) > limit:
                    raise FlowError(E.INVALID_REQUEST, "Image exceeds 20 MiB")
                return data
            finally:
                connection.close()
        except (OSError, ValueError):
            raise FlowError(E.INVALID_REQUEST, "Unable to fetch image URL") from None