import os
import sqlite3
import time
import uuid
from datetime import datetime
import cv2
import numpy as np
from flask import (
    Flask,
    flash,
    jsonify,
    redirect,
    render_template,
    request,
    send_file,
    session,
    url_for,
)
from flask_login import LoginManager, UserMixin, current_user, login_required, login_user, logout_user
from werkzeug.security import check_password_hash, generate_password_hash
from werkzeug.utils import secure_filename

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
UPLOAD_FOLDER = os.path.join(BASE_DIR, "static", "uploads")
RESULT_FOLDER = os.path.join(BASE_DIR, "static", "results")
DATABASE = os.path.join(BASE_DIR, "database.db")

ALLOWED_EXTENSIONS = {"png", "jpg", "jpeg", "webp"}
MAX_CONTENT_LENGTH = 16 * 1024 * 1024  # 16 MB

app = Flask(__name__)
app.config["SECRET_KEY"] = os.environ.get("PIXELSAGE_SECRET", "pixelsage-dev-secret-change-in-production")
app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER
app.config["RESULT_FOLDER"] = RESULT_FOLDER
app.config["MAX_CONTENT_LENGTH"] = MAX_CONTENT_LENGTH

os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(RESULT_FOLDER, exist_ok=True)

login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = "login"
login_manager.login_message_category = "info"


# ---------------------------------------------------------------------------
# Database
# ---------------------------------------------------------------------------

def get_db():
    conn = sqlite3.connect(DATABASE)
    conn.row_factory = sqlite3.Row
    return conn


def init_db():
    with get_db() as conn:
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                password_hash TEXT NOT NULL,
                created_at TEXT NOT NULL
            );
            CREATE TABLE IF NOT EXISTS history (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                operation_name TEXT NOT NULL,
                original_filename TEXT NOT NULL,
                result_filename TEXT,
                processing_time REAL,
                created_at TEXT NOT NULL,
                FOREIGN KEY (user_id) REFERENCES users (id)
            );
        """)


# ---------------------------------------------------------------------------
# User model
# ---------------------------------------------------------------------------

class User(UserMixin):
    def __init__(self, id_, username):
        self.id = id_
        self.username = username


@login_manager.user_loader
def load_user(user_id):
    with get_db() as conn:
        row = conn.execute("SELECT id, username FROM users WHERE id = ?", (user_id,)).fetchone()
    if row:
        return User(row["id"], row["username"])
    return None


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def allowed_file(filename):
    return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS


def read_image_from_bytes(data):
    arr = np.frombuffer(data, dtype=np.uint8)
    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError("Could not decode image. Please upload a valid JPG or PNG file.")
    return img


def encode_image(img, ext=".png"):
    ext = ext.lower()
    if ext in (".jpg", ".jpeg"):
        _, buf = cv2.imencode(".jpg", img, [int(cv2.IMWRITE_JPEG_QUALITY), 92])
    else:
        _, buf = cv2.imencode(".png", img)
    return buf.tobytes()


def save_history(user_id, operation_name, original_filename, result_filename, processing_time):
    with get_db() as conn:
        conn.execute(
            """INSERT INTO history
               (user_id, operation_name, original_filename, result_filename, processing_time, created_at)
               VALUES (?, ?, ?, ?, ?, ?)""",
            (
                user_id,
                operation_name,
                original_filename,
                result_filename,
                round(processing_time, 4),
                datetime.utcnow().isoformat(),
            ),
        )


def get_user_history(user_id, limit=10):
    with get_db() as conn:
        rows = conn.execute(
            """SELECT operation_name, original_filename, processing_time, created_at
               FROM history WHERE user_id = ? ORDER BY id DESC LIMIT ?""",
            (user_id, limit),
        ).fetchall()
    return [dict(r) for r in rows]


# ---------------------------------------------------------------------------
# Image operations
# ---------------------------------------------------------------------------

OPERATIONS = {
    "sobel": "Sobel Edge Detection",
    "affine": "Affine Transform",
    "gaussian_blur": "Gaussian Blur",
    "sharpen": "Image Sharpening",
    "grayscale": "Grayscale",
    "histogram_eq": "Histogram Equalization",
}


def apply_sobel(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    grad_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
    grad_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
    magnitude = cv2.magnitude(grad_x, grad_y)
    magnitude = cv2.normalize(magnitude, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
    return cv2.cvtColor(magnitude, cv2.COLOR_GRAY2BGR)


def apply_affine(img, angle=15, scale=1.0, tx=0, ty=0):
    h, w = img.shape[:2]
    center = (w // 2, h // 2)
    matrix = cv2.getRotationMatrix2D(center, float(angle), float(scale))
    matrix[0, 2] += float(tx)
    matrix[1, 2] += float(ty)
    return cv2.warpAffine(img, matrix, (w, h), borderMode=cv2.BORDER_REFLECT)


def apply_gaussian_blur(img, ksize=15):
    k = max(3, int(ksize) | 1)  # ensure odd >= 3
    return cv2.GaussianBlur(img, (k, k), 0)


def apply_sharpen(img):
    kernel = np.array(
        [[0, -1, 0], [-1, 5, -1], [0, -1, 0]], dtype=np.float32
    )
    return cv2.filter2D(img, -1, kernel)


def apply_grayscale(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)


def apply_histogram_eq(img):
    ycrcb = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
    ycrcb[:, :, 0] = cv2.equalizeHist(ycrcb[:, :, 0])
    return cv2.cvtColor(ycrcb, cv2.COLOR_YCrCb2BGR)


def process_image(img, operation, params=None):
    params = params or {}
    ops = {
        "sobel": lambda i: apply_sobel(i),
        "affine": lambda i: apply_affine(
            i,
            angle=params.get("angle", 15),
            scale=params.get("scale", 1.0),
            tx=params.get("tx", 0),
            ty=params.get("ty", 0),
        ),
        "gaussian_blur": lambda i: apply_gaussian_blur(i, ksize=params.get("blur_ksize", 15)),
        "sharpen": lambda i: apply_sharpen(i),
        "grayscale": lambda i: apply_grayscale(i),
        "histogram_eq": lambda i: apply_histogram_eq(i),
    }
    if operation not in ops:
        raise ValueError(f"Unknown operation: {operation}")
    return ops[operation](img)


def parse_affine_params(form):
    def f(key, default):
        try:
            return float(form.get(key, default))
        except (TypeError, ValueError):
            return default

    return {
        "angle": f("angle", 15),
        "scale": f("scale", 1.0),
        "tx": f("tx", 0),
        "ty": f("ty", 0),
    }


def parse_blur_params(form):
    try:
        k = int(form.get("blur_ksize", 15))
    except (TypeError, ValueError):
        k = 15
    return {"blur_ksize": max(3, min(k, 51))}


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@app.route("/")
def index():
    return render_template("index.html")


@app.route("/register", methods=["GET", "POST"])
def register():
    if current_user.is_authenticated:
        return redirect(url_for("dashboard"))

    if request.method == "POST":
        username = (request.form.get("username") or "").strip()
        password = request.form.get("password") or ""
        confirm = request.form.get("confirm_password") or ""

        if not username or not password:
            flash("Username and password are required.", "error")
        elif len(username) < 3:
            flash("Username must be at least 3 characters.", "error")
        elif len(password) < 6:
            flash("Password must be at least 6 characters.", "error")
        elif password != confirm:
            flash("Passwords do not match.", "error")
        else:
            try:
                with get_db() as conn:
                    conn.execute(
                        "INSERT INTO users (username, password_hash, created_at) VALUES (?, ?, ?)",
                        (username, generate_password_hash(password), datetime.utcnow().isoformat()),
                    )
                flash("Account created successfully. Please log in.", "success")
                return redirect(url_for("login"))
            except sqlite3.IntegrityError:
                flash("Username already taken. Choose another.", "error")

    return render_template("register.html")


@app.route("/login", methods=["GET", "POST"])
def login():
    if current_user.is_authenticated:
        return redirect(url_for("dashboard"))

    if request.method == "POST":
        username = (request.form.get("username") or "").strip()
        password = request.form.get("password") or ""

        with get_db() as conn:
            row = conn.execute(
                "SELECT id, username, password_hash FROM users WHERE username = ?", (username,)
            ).fetchone()

        if row and check_password_hash(row["password_hash"], password):
            login_user(User(row["id"], row["username"]), remember=True)
            flash(f"Welcome back, {row['username']}!", "success")
            next_page = request.args.get("next")
            return redirect(next_page or url_for("dashboard"))
        flash("Invalid username or password.", "error")

    return render_template("login.html")


@app.route("/logout")
@login_required
def logout():
    logout_user()
    flash("You have been logged out.", "info")
    return redirect(url_for("index"))


@app.route("/dashboard")
@login_required
def dashboard():
    history = get_user_history(current_user.id)
    return render_template("dashboard.html", operations=OPERATIONS, history=history)


@app.route("/api/process", methods=["POST"])
@login_required
def api_process():
    if "image" not in request.files:
        return jsonify({"error": "No image uploaded."}), 400

    file = request.files["image"]
    operation = request.form.get("operation", "sobel")

    if file.filename == "":
        return jsonify({"error": "No file selected."}), 400
    if not allowed_file(file.filename):
        return jsonify({"error": "Invalid file type. Use JPG or PNG."}), 400
    if operation not in OPERATIONS:
        return jsonify({"error": "Invalid operation selected."}), 400

    params = {}
    if operation == "affine":
        params = parse_affine_params(request.form)
    elif operation == "gaussian_blur":
        params = parse_blur_params(request.form)

    try:
        raw = file.read()
        img = read_image_from_bytes(raw)
        start = time.perf_counter()
        result = process_image(img, operation, params)
        elapsed = time.perf_counter() - start

        ext = os.path.splitext(secure_filename(file.filename))[1] or ".png"
        result_name = f"{uuid.uuid4().hex}{ext}"
        result_path = os.path.join(app.config["RESULT_FOLDER"], result_name)
        with open(result_path, "wb") as f:
            f.write(encode_image(result, ext))

        original_name = secure_filename(file.filename)
        save_history(
            current_user.id,
            OPERATIONS[operation],
            original_name,
            result_name,
            elapsed,
        )

        session["last_result"] = result_name
        session["last_operation"] = OPERATIONS[operation]

        return jsonify(
            {
                "success": True,
                "result_url": url_for("static", filename=f"results/{result_name}"),
                "download_url": url_for("download_result", filename=result_name),
                "processing_time": round(elapsed, 4),
                "operation": OPERATIONS[operation],
            }
        )
    except ValueError as e:
        return jsonify({"error": str(e)}), 400
    except Exception:
        return jsonify({"error": "Processing failed. Please try another image."}), 500


@app.route("/download/<filename>")
@login_required
def download_result(filename):
    safe = secure_filename(filename)
    path = os.path.join(app.config["RESULT_FOLDER"], safe)
    if not os.path.isfile(path):
        flash("File not found or expired.", "error")
        return redirect(url_for("dashboard"))
    return send_file(path, as_attachment=True, download_name=f"pixelsage_{safe}")


@app.errorhandler(413)
def too_large(_):
    flash("File is too large. Maximum size is 16 MB.", "error")
    return redirect(url_for("dashboard"))


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    init_db()
    app.run(debug=True, host="0.0.0.0", port=5000)
