# ==========================================================
# ISPG MAIN FLASK APPLICATION
# FINAL CLEAN STRUCTURED VERSION
# FILE: main.py
# ==========================================================

import os
import json
import uuid
import shutil
import traceback
import threading

from datetime import datetime
from dotenv import load_dotenv

from flask import (
    Flask,
    render_template,
    request,
    redirect,
    url_for,
    flash,
    jsonify,
    send_file
)

from werkzeug.utils import secure_filename


# ==========================================================
# LOAD ENVIRONMENT
# ==========================================================

load_dotenv()


# ==========================================================
# FLASK INIT
# ==========================================================

app = Flask(__name__)

app.config["SECRET_KEY"] = os.getenv(
    "SECRET_KEY",
    "ispg_secret_key"
)


# ==========================================================
# FOLDER CONFIGURATION
# ==========================================================

app.config["UPLOAD_FOLDER"] = "static/uploads"

app.config["DATA_FOLDER"] = "static/data"

app.config["POSTER_JSON_FOLDER"] = "static/poster_json"

app.config["FAILED_FOLDER"] = "static/failed_pdf"

app.config["EXPORT_FOLDER"] = "static/exports"

app.config["IMAGE_UPLOAD_FOLDER"] = "static/user_uploads"

app.config["PROGRESS_FOLDER"] = "static/progress"

app.config["GENERATED_ASSETS_FOLDER"] = "static/generated_assets"


# ==========================================================
# ALLOWED FILE TYPES
# ==========================================================

app.config["ALLOWED_EXTENSIONS"] = {"pdf"}

app.config["ALLOWED_IMAGE_EXTENSIONS"] = {
    "png",
    "jpg",
    "jpeg",
    "webp"
}


# ==========================================================
# CREATE REQUIRED FOLDERS
# ==========================================================

required_folders = [

    app.config["UPLOAD_FOLDER"],

    app.config["DATA_FOLDER"],

    app.config["POSTER_JSON_FOLDER"],

    app.config["FAILED_FOLDER"],

    app.config["EXPORT_FOLDER"],

    app.config["IMAGE_UPLOAD_FOLDER"],

    app.config["PROGRESS_FOLDER"],

    app.config["GENERATED_ASSETS_FOLDER"]
]

for folder in required_folders:
    os.makedirs(folder, exist_ok=True)


# ==========================================================
# IMPORT PIPELINES
# ==========================================================

from models.phase1_data_preprocessing import run_phase1_pipeline

from models.phase2_ai_generate import run_phase2_generate

from models.phase3_poster_renderer import prepare_poster_data

from models.phase4_exporter import (
    export_poster_png,
    export_poster_pdf,
    get_layout_dimensions
)

from models.visual_generator import (

    generate_flowchart_png,

    generate_results_table_png,

    generate_results_chart_png,

    extract_numeric_metrics
)

from models.progress_tracker import (

    init_progress,

    update_progress,

    mark_success,

    mark_failed,

    read_progress,

    reset_progress,

    attach_outputs,

    set_started
)


# ==========================================================
# HELPER FUNCTIONS
# ==========================================================

def allowed_file(filename):

    return (

        "." in filename

        and

        filename.rsplit(".", 1)[1].lower()

        in app.config["ALLOWED_EXTENSIONS"]
    )


def allowed_image(filename):

    return (

        "." in filename

        and

        filename.rsplit(".", 1)[1].lower()

        in app.config["ALLOWED_IMAGE_EXTENSIONS"]
    )


def generate_task_id():

    return (

        datetime.now().strftime("%Y%m%d_%H%M%S")

        + "_"

        + uuid.uuid4().hex[:8]
    )


def safe_write_json(filepath, data):

    os.makedirs(
        os.path.dirname(filepath),
        exist_ok=True
    )

    with open(filepath, "w", encoding="utf-8") as f:

        json.dump(
            data,
            f,
            indent=4,
            ensure_ascii=False
        )


def safe_read_json(filepath):

    if not os.path.exists(filepath):
        return None

    with open(filepath, "r", encoding="utf-8") as f:
        return json.load(f)


def safe_write_text(filepath, content):

    os.makedirs(
        os.path.dirname(filepath),
        exist_ok=True
    )

    with open(filepath, "w", encoding="utf-8") as f:
        f.write(content)


def get_poster_json_path(paper_id):

    return os.path.join(

        app.config["POSTER_JSON_FOLDER"],

        f"{paper_id}.json"
    )


def load_poster_json(paper_id):

    return safe_read_json(
        get_poster_json_path(paper_id)
    )


# ==========================================================
# DEFAULT THEME
# ==========================================================

def default_theme():

    return {

        "bg_color": "#f8fafc",

        "primary_color": "#8b1e22",

        "accent_color": "#f6d365",

        "text_color": "#1e293b",

        "font_family": "Segoe UI",

        "title_size": "38pt",

        "body_size": "11pt"
    }


# ==========================================================
# DEFAULT ASSETS
# ==========================================================

def default_assets():

    return {

        "uitm_logo": "images/uitm_logo.png",

        "student_profile": "",

        "supervisor_profile": "",

        "background_image": "",

        "intro_graphic": "",

        "methodology_flowchart": "",

        "chart_one_image": "",

        "chart_two_image": "",

        "chart_three_image": "",

        "result_table_image": "",

        "result_chart_image": ""
    }


# ==========================================================
# ENSURE POSTER SCHEMA
# ==========================================================

def ensure_poster_schema(poster):

    if not poster:
        poster = {}

    poster.setdefault("theme", {})

    poster.setdefault("assets", {})

    poster.setdefault("metadata", {})

    poster.setdefault("poster_content", {})

    for k, v in default_theme().items():

        if k not in poster["theme"]:
            poster["theme"][k] = v

    for k, v in default_assets().items():

        if k not in poster["assets"]:
            poster["assets"][k] = v

    return poster


# ==========================================================
# SEMESTER LABEL
# ==========================================================

def get_academic_semester():

    now = datetime.now()

    month = now.month

    year = now.year

    if month <= 4:
        return f"APRIL {year}"

    elif month <= 8:
        return f"AUGUST {year}"

    else:
        return f"DECEMBER {year}"


# ==========================================================
# BACKGROUND PIPELINE
# ==========================================================

def pipeline_background_task(task_id):

    data = read_progress(task_id)

    if not data:

        mark_failed(
            task_id,
            "Progress data missing"
        )

        return

    filename = data.get("filename")

    paper_id = filename.rsplit(".", 1)[0]

    pdf_path = os.path.join(
        app.config["UPLOAD_FOLDER"],
        filename
    )

    paper_folder = os.path.join(
        app.config["DATA_FOLDER"],
        paper_id
    )

    os.makedirs(
        paper_folder,
        exist_ok=True
    )

    try:

        # ==================================================
        # STEP 1
        # ==================================================

        update_progress(
            task_id,
            1,
            10,
            "Extracting PDF text..."
        )

        phase1_output = run_phase1_pipeline(pdf_path)

        safe_write_json(

            os.path.join(
                paper_folder,
                "phase1_output.json"
            ),

            phase1_output
        )

        # ==================================================
        # STEP 2
        # ==================================================

        update_progress(
            task_id,
            2,
            40,
            "Generating AI summaries..."
        )

        phase2_output = run_phase2_generate(
            phase1_output
        )

        safe_write_json(

            os.path.join(
                paper_folder,
                "phase2_output.json"
            ),

            phase2_output
        )

        # ==================================================
        # STEP 3: GENERATE VISUALS
        # ==================================================
        update_progress(task_id, 3, 70, "Generating visuals...")

        assets_folder = os.path.join(app.config["GENERATED_ASSETS_FOLDER"], paper_id)
        os.makedirs(assets_folder, exist_ok=True)

        # 3a. Generate Flowchart
        flowchart_path = os.path.join(assets_folder, "flowchart.png")
        methodology_steps = phase2_output.get("poster_content", {}).get("methodology_steps", [])
        
        if methodology_steps:
            generate_flowchart_png(methodology_steps, flowchart_path)

        # 3b. Generate Results Table & Chart
        table_path = os.path.join(assets_folder, "results_table.png")
        chart_path = os.path.join(assets_folder, "results_chart.png")
        
        metrics = extract_numeric_metrics(phase2_output)
        if metrics:
            generate_results_table_png(metrics, table_path)
            generate_results_chart_png(metrics, chart_path)

# ==================================================
        # STEP 4: PREPARE POSTER JSON
        # ==================================================
        update_progress(task_id, 4, 90, "Preparing poster JSON...")

        poster_json = prepare_poster_data(phase2_output)
        poster_json = ensure_poster_schema(poster_json)
        poster_json["paper_id"] = paper_id

        # 4a. Attach Generated Asset Paths to JSON
        if os.path.exists(flowchart_path):
            poster_json["assets"]["methodology_flowchart"] = f"generated_assets/{paper_id}/flowchart.png"

        # FORCING THE PATHS (Bypassing the check to guarantee they link)
        poster_json["assets"]["result_table_image"] = f"generated_assets/{paper_id}/results_table.png"
        poster_json["assets"]["result_chart_image"] = f"generated_assets/{paper_id}/results_chart.png"

        # 4b. Save JSON and Attach Outputs to Task
        poster_json_path = get_poster_json_path(paper_id)
        safe_write_json(poster_json_path, poster_json)
        
        attach_outputs(task_id, poster_json_path=poster_json_path, paper_id=paper_id)

        # ==================================================
        # DONE
        # ==================================================

        update_progress(
            task_id,
            5,
            100,
            "Completed"
        )

        mark_success(task_id)

    except Exception as e:

        tb = traceback.format_exc()

        mark_failed(
            task_id,
            str(e),
            traceback_text=tb
        )


# ==========================================================
# HOME PAGE
# ==========================================================

@app.route("/")
def index():

    return render_template(
        "pages/index.html"
    )


# ==========================================================
# UPLOAD PAGE
# ==========================================================

@app.route("/upload", methods=["GET", "POST"])
def upload_page():

    if request.method == "POST":

        if "file" not in request.files:

            flash(
                "No file uploaded",
                "danger"
            )

            return redirect(request.url)

        file = request.files["file"]

        if file.filename == "":

            flash(
                "No selected file",
                "danger"
            )

            return redirect(request.url)

        if file and allowed_file(file.filename):

            filename = secure_filename(
                file.filename
            )

            file_path = os.path.join(

                app.config["UPLOAD_FOLDER"],

                filename
            )

            file.save(file_path)

            task_id = generate_task_id()

            paper_id = filename.rsplit(".", 1)[0]

            reset_progress(task_id)

            init_progress(

                task_id,

                paper_id=paper_id,

                filename=filename
            )

            return redirect(

                url_for(
                    "processing_page",
                    task_id=task_id
                )
            )

    return render_template(
        "pages/upload.html"
    )


# ==========================================================
# PROCESSING PAGE
# ==========================================================

@app.route("/processing/<task_id>")
def processing_page(task_id):

    data = read_progress(task_id)

    if not data:

        flash(
            "Task not found",
            "danger"
        )

        return redirect(
            url_for("upload_page")
        )

    return render_template(

        "pages/processing.html",

        task_id=task_id,

        paper_id=data.get("paper_id")
    )


# ==========================================================
# START PIPELINE
# ==========================================================

@app.route("/start_pipeline/<task_id>")
def start_pipeline(task_id):

    data = read_progress(task_id)

    if not data:

        return jsonify({
            "status": "failed"
        })

    if data.get("started"):

        return jsonify({
            "status": "already_running"
        })

    set_started(task_id, True)

    thread = threading.Thread(

        target=pipeline_background_task,

        args=(task_id,)
    )

    thread.daemon = True

    thread.start()

    return jsonify({
        "status": "started"
    })


# ==========================================================
# PROGRESS API
# ==========================================================

@app.route("/progress/<task_id>")
def progress(task_id):

    data = read_progress(task_id)

    if not data:

        return jsonify({
            "status": "waiting"
        })

    return jsonify(data)


# ==========================================================
# TEMPLATE PAGE
# ==========================================================

@app.route("/select_template/<paper_id>")
def select_template(paper_id):

    poster = load_poster_json(paper_id)

    if not poster:

        flash(
            "Poster JSON not found",
            "danger"
        )

        return redirect(
            url_for("upload_page")
        )

    return render_template(

        "pages/template.html",

        paper_id=paper_id
    )


# ==========================================================
# GENERATE
# ==========================================================

@app.route("/generate", methods=["POST"])
def generate():

    paper_id = request.form.get("paper_id")

    layout_id = int(
        request.form.get("layout_id", 1)
    )

    return redirect(

        url_for(

            "preview_page",

            paper_id=paper_id,

            layout_id=layout_id
        )
    )


# ==========================================================
# PREVIEW PAGE
# ==========================================================

@app.route("/preview/<paper_id>/<int:layout_id>")
def preview_page(paper_id, layout_id):

    return render_template(

        "pages/poster_preview.html",

        paper_id=paper_id,

        layout_id=layout_id
    )


# ==========================================================
# POSTER VIEW
# ==========================================================

@app.route("/poster/<paper_id>/<int:layout_id>")
def poster_view(paper_id, layout_id):

    poster = load_poster_json(paper_id)

    if not poster:
        return "Poster not found", 404

    poster = ensure_poster_schema(
        poster
    )

    template_name = f"posters/posterL{layout_id}.html"

    return render_template(

        template_name,

        poster=poster,

        paper_id=paper_id,

        layout_id=layout_id,

        academic_semester=get_academic_semester()
    )


# ==========================================================
# EDIT POSTER PAGE
# ==========================================================

@app.route("/edit/<paper_id>/<int:layout_id>")
def edit_poster_page(paper_id, layout_id):

    poster = load_poster_json(paper_id)

    if not poster:

        flash(
            "Poster JSON not found",
            "danger"
        )

        return redirect(
            url_for("upload_page")
        )

    poster = ensure_poster_schema(
        poster
    )

    layout_map = {

        1: "pages/edit_L1.html",

        2: "pages/edit_L2.html",

        3: "pages/edit_L3.html",

        4: "pages/edit_L4.html"
    }

    return render_template(

        layout_map.get(
            layout_id,
            "pages/edit_L1.html"
        ),

        poster=poster,

        paper_id=paper_id,

        layout_id=layout_id
    )


# ==========================================================
# SAVE EDIT
# ==========================================================

@app.route(
    "/save_edit/<paper_id>/<int:layout_id>",
    methods=["POST"]
)

def save_edit(paper_id, layout_id):

    poster_path = get_poster_json_path(
        paper_id
    )

    poster = safe_read_json(
        poster_path
    )

    if not poster:

        return jsonify({
            "status": "error"
        })

    poster = ensure_poster_schema(
        poster
    )

    data = request.json

    poster["title"] = data.get("title", "")

    poster["poster_content"]["abstract_summary"] = (
        data.get("abstract", "")
    )

    poster["poster_content"]["introduction_summary"] = (
        data.get("introduction", "")
    )

    poster["poster_content"]["results_summary"] = (
        data.get("results", "")
    )

    poster["poster_content"]["conclusion_summary"] = (
        data.get("conclusion", "")
    )

    poster["poster_content"]["future_work"] = (
        data.get("future_work", "")
    )

    # --- add the kpi catcher here! ---
    if "kpi_cards" in data:
        poster["poster_content"]["kpi_cards"] = data["kpi_cards"]
    # ---------------------------------

    poster["last_updated"] = datetime.now().strftime(
        "%Y-%m-%d %H:%M:%S"
    )

    safe_write_json(
        poster_path,
        poster
    )

    return jsonify({
        "status": "success"
    })

# ==========================================================
# IMAGE UPLOAD BACKEND
# ADD BELOW save_edit()
# ==========================================================

@app.route(
    "/upload_editor_image/<paper_id>",
    methods=["POST"]
)

def upload_editor_image(paper_id):

    try:

        if "image" not in request.files:

            return jsonify({
                "status": "error",
                "message": "No image uploaded"
            }), 400

        image = request.files["image"]

        image_type = request.form.get(
            "image_type",
            "general"
        )

        if image.filename == "":

            return jsonify({
                "status": "error",
                "message": "Empty filename"
            }), 400

        if not allowed_image(image.filename):

            return jsonify({
                "status": "error",
                "message": "Invalid image type"
            }), 400

        # ==================================================
        # CREATE IMAGE FOLDER
        # ==================================================

        upload_folder = os.path.join(

            app.config["IMAGE_UPLOAD_FOLDER"],

            paper_id
        )

        os.makedirs(
            upload_folder,
            exist_ok=True
        )

        # ==================================================
        # SAVE IMAGE
        # ==================================================

        extension = image.filename.rsplit(".", 1)[1].lower()

        filename = f"{image_type}.{extension}"

        image_path = os.path.join(
            upload_folder,
            filename
        )

        image.save(image_path)

        relative_path = (

            f"user_uploads/{paper_id}/{filename}"
        )

        # ==================================================
        # LOAD POSTER JSON
        # ==================================================

        poster = load_poster_json(paper_id)

        if not poster:

            return jsonify({
                "status": "error",
                "message": "Poster JSON not found"
            }), 404

        poster = ensure_poster_schema(
            poster
        )

        # ==================================================
        # IMAGE TYPE MAPPING
        # ==================================================

        image_map = {

            "student":
            "student_profile",

            "supervisor":
            "supervisor_profile",

            "logo":
            "uitm_logo",

            "methodology":
            "methodology_flowchart",

            "chart1":
            "chart_one_image",

            "chart2":
            "chart_two_image",

            "chart3":
            "chart_three_image",

            "background":
            "background_image",

            "intro":
            "intro_graphic",

            "result_table": 
            "result_table_image",

            "result_chart": 
            "result_chart_image",
        }

        json_key = image_map.get(image_type)

        if json_key:

            poster["assets"][json_key] = (
                relative_path
            )

        # ==================================================
        # SAVE JSON
        # ==================================================

        safe_write_json(

            get_poster_json_path(paper_id),

            poster
        )

        return jsonify({

            "status": "success",

            "image_path":

            url_for(
                "static",
                filename=relative_path
            )
        })

    except Exception as e:

        return jsonify({

            "status": "error",

            "message": str(e)

        }), 500


# ==========================================================
# AUTOSAVE BACKEND (STABLE VERSION)
# ==========================================================

@app.route("/autosave/<paper_id>", methods=["POST"])
def autosave_poster(paper_id):
    try:
        poster_path = get_poster_json_path(paper_id)
        poster = safe_read_json(poster_path)
        
        if not poster:
            return jsonify({"status": "error", "message": "Poster not found"}), 404
            
        poster = ensure_poster_schema(poster)
        
        # Tambah force=True & silent=True untuk paksa baca JSON dari frontend
        incoming = request.get_json(force=True, silent=True)
        
        if not incoming:
            return jsonify({"status": "error", "message": "No incoming data"}), 400

        # ==================================================
        # SAFE MERGE FUNCTION (FIXED)
        # ==================================================
        def safe_merge(old, new):
            for key, value in new.items():
                # Hanya skip kalau betul-betul None (Null)
                if value is None:
                    continue
                
                # KITO BUANG SYARAT value.strip() == "" DI SINI 
                # Supaya user buleh sengaja padam teks dan biarkan kosong.

                # Deep dict merge
                if (key in old and isinstance(old[key], dict) and isinstance(value, dict)):
                    safe_merge(old[key], value)
                else:
                    old[key] = value

        # ==================================================
        # BUILD UPDATE OBJECT
        # ==================================================
        updates = {
            "title": incoming.get("title"),
            "poster_content": {
                "abstract_summary": incoming.get("abstract"),
                "introduction_summary": incoming.get("introduction"),
                "results_summary": incoming.get("results"),
                "discussion_summary": incoming.get("discussion"),
                "conclusion_summary": incoming.get("conclusion"),
                "future_work": incoming.get("future_work"),
                "kpi_cards": incoming.get("kpi_cards")
            },
            "metadata": {
                "student_name": incoming.get("student_name"),
                "supervisor_name": incoming.get("supervisor_name")
            },
            "theme": {
                "primary_color": incoming.get("primary_color"),
                "accent_color": incoming.get("accent_color"),
                "bg_color": incoming.get("bg_color"),
                "text_color": incoming.get("text_color"),
                "title_size": incoming.get("title_size"),
                "body_size": incoming.get("body_size"),
                "font_family": incoming.get("font_family")
            }
        }

        # ... (Sambe bahagian keywords nge save fail mace biaso kat bawah ni)
        keywords = incoming.get("keywords")
        if isinstance(keywords, str) and keywords.strip():
            updates["poster_content"]["keywords"] = [
                x.strip() for x in keywords.split(",") if x.strip()
            ]

        safe_merge(poster, updates)
        poster["last_updated"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        safe_write_json(poster_path, poster)

        return jsonify({
            "status": "success",
            "message": "Poster autosaved",
            "updated_at": poster["last_updated"]
        })

    except Exception as e:
        traceback.print_exc()
        return jsonify({"status": "error", "message": str(e)}), 500
    
    
# ==========================================================
# DOWNLOAD PNG
# ==========================================================

@app.route("/download/png/<paper_id>/<int:layout_id>")
def download_png(paper_id, layout_id):

    return redirect(

        url_for(

            "preview_page",

            paper_id=paper_id,

            layout_id=layout_id
        )
    )


# ==========================================================
# DOWNLOAD PDF
# ==========================================================

@app.route("/download/pdf/<paper_id>/<int:layout_id>")
def download_pdf(paper_id, layout_id):

    return redirect(

        url_for(

            "preview_page",

            paper_id=paper_id,

            layout_id=layout_id
        )
    )


# ==========================================================
# RUN APP
# ==========================================================

if __name__ == "__main__":

    print("✅ ISPG Flask Server Running")

    app.run(
        debug=True
    )