from __future__ import annotations

import base64
import csv
import io
import mimetypes
import zipfile
from dataclasses import dataclass
from typing import Any, List, Tuple

import dataiku
from flask import Blueprint, abort, g, send_file

from backend.config import get_admin_uploads_managedfolder_id, get_uploads_managedfolder_id
from backend.constants import DIP_MIMETYPE_EXTENSIONS, PDF_MIME_TYPE
from backend.schemas import schemas
from backend.services.generated_document.generated_document_metadata import is_history_marker
from backend.services.generated_document.generated_document_preview_service import (
    PreviewArtifactGone,
    PreviewArtifactNotFound,
    PreviewConversionError,
    build_generated_document_preview_pdf,
)
from backend.utils.logger_utils import log_http_request
from backend.utils.logging_utils import get_logger
from backend.utils.utils import get_store

logger = get_logger(__name__)

downloads_bp = Blueprint("downloads", __name__, url_prefix="/downloads")


@dataclass
class Table:
    columns: list[str]
    data: list[list[Any]]


def sanitize_spreadsheet_cell(value: Any) -> Any:
    """Prevent strings from being interpreted as formulas by spreadsheet clients."""
    if not isinstance(value, str) or not value:
        return value

    for character in value:
        if character in "\t\r":
            return f"'{value}"
        if character.isspace():
            continue
        if character in "=+-@":
            return f"'{value}"
        break

    return value


def to_csv_string(table: Table) -> str:
    """Convert a Table {columns, data} into a CSV string."""
    if not table or not isinstance(table.columns, list) or not isinstance(table.data, list):
        raise ValueError("Event data is not in { columns, data } format")

    output = io.StringIO()
    writer = csv.writer(output, quoting=csv.QUOTE_MINIMAL)
    writer.writerow(sanitize_spreadsheet_cell(column) for column in table.columns)
    for row in table.data:
        writer.writerow(sanitize_spreadsheet_cell(value) for value in row)
    return output.getvalue()


def records_folder_from_artifacts(artifacts: list[schemas.Artifact]) -> dict[str, list[Table]]:
    """
    Convert artifacts to { folderName: [Table, Table, ...] }.
    """
    folders: dict[str, list[Table]] = {}
    for index, art in enumerate(artifacts) or []:
        for item in art.get("parts") or []:
            if item.get("type") == "RECORDS":
                if art["name"] + str(index) not in folders:
                    folders[art["name"] + str(index)] = []
                columns = item.get("records", {}).get("columns", None)
                data = item.get("records", {}).get("data", None)
                if (columns is not None) and (data is not None):
                    folders[art["name"] + str(index)].append(Table(columns=columns, data=data))
    return folders


def extract_recs(artifact: dict[str, Any]) -> list[Table]:
    recs: list[Table] = []
    for item in artifact.get("parts") or []:
        if item.get("type") != "RECORDS":
            continue
        columns = item.get("records", {}).get("columns")
        data = item.get("records", {}).get("data")
        if columns is not None and data is not None:
            recs.append(Table(columns=columns, data=data))
    return recs


def build_csv_files_from_records(records: list[Table], info: str) -> list[Tuple[str, str]]:
    """
    For a list of tables, build CSV strings.
    Returns list of (filename, csv_string).
    """
    files: list[Tuple[str, str]] = []
    for idx, table in enumerate(records):
        csv_str = to_csv_string(table)
        safe_info = (info or "artifact").strip().replace(" ", "_")
        files.append((f"{safe_info}_{idx}.csv", csv_str))
    return files


def _get_art_grp(msg_id: str, art_id: str) -> dict[str, Any]:
    store = get_store()
    artifacts_meta = store.get_message_artifacts_meta(msg_id)
    if not isinstance(artifacts_meta, dict) or art_id not in artifacts_meta:
        abort(404, description=f"Artifact id '{art_id}' not found for message '{msg_id}'")
    art_grp = artifacts_meta[art_id]
    if not isinstance(art_grp, dict) or not isinstance(art_grp.get("artifacts"), list):
        abort(404, description=f"Artifact id '{art_id}' is malformed for message '{msg_id}'")
    return art_grp


def _get_art(msg_id: str, art_id: str, artifact_index: str) -> tuple[int, dict[str, Any], dict[str, Any]]:
    try:
        art_idx = int(artifact_index)
    except ValueError:
        abort(404, description=f"Artifact index '{artifact_index}' should be a number for artifact id '{art_id}'")

    art_grp = _get_art_grp(msg_id, art_id)
    artifacts = art_grp["artifacts"]
    if art_idx < 0 or art_idx >= len(artifacts):
        abort(404, description=f"Artifact index '{artifact_index}' out of range for artifact id '{art_id}'")
    art = artifacts[art_idx]
    if not isinstance(art, dict):
        abort(404, description=f"Artifact index '{artifact_index}' is malformed for artifact id '{art_id}'")
    if is_history_marker(art):
        abort(410, description="Generated document content is no longer retained")
    return art_idx, art_grp, art


def _get_rec_tbl(
    *, msg_id: str, art_id: str, artifact_index: str, records_index: str
) -> tuple[dict[str, Any], dict[str, Any], Table]:
    try:
        rec_idx = int(records_index)
    except ValueError:
        abort(404, description=f"Records index '{records_index}' should be a number for artifact id '{art_id}'")

    _, art_grp, art = _get_art(msg_id, art_id, artifact_index)
    recs = extract_recs(art)
    if rec_idx < 0 or rec_idx >= len(recs):
        abort(
            404,
            description=(
                f"Records index '{records_index}' out of range for artifact id '{art_id}' "
                f"and artifact index '{artifact_index}'"
            ),
        )
    return art_grp, art, recs[rec_idx]


def _safe(s: str) -> str:
    return (s or "").strip().replace(" ", "_").replace("/", "_")


def build_zip_bytes(aggregated_folders: dict[str, dict[str, list]], inline_data_files: list[dict]) -> bytes:
    """
    folders: { folderName: [(filename, csv_str), ...], ... }
    inline_data_files: [{mimeType:'application/pdf', 'dataBase64': 'Qkx42'}, ...]
    Returns zip bytes.
    """

    mem = io.BytesIO()
    with zipfile.ZipFile(mem, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        for source, folders in aggregated_folders.items():
            source_prefix = _safe(source or "exports")
            for folder, files in folders.items():
                folder_prefix = f"{source_prefix}/{_safe(folder)}/"
                for filename, csv_str in files:
                    zf.writestr(folder_prefix + filename, csv_str)

        for i, item in enumerate(inline_data_files):
            mimeType = item.get("mimeType")
            dataBase64 = item.get("dataBase64")
            filename = item.get("filename", None)
            if not filename:
                extension = DIP_MIMETYPE_EXTENSIONS.get(mimeType, "bin")
                filename = f"document_{i}.{extension}"

            if mimeType and dataBase64:
                try:
                    decoded_bytes = base64.b64decode(dataBase64)
                    zf.writestr(_safe(filename), decoded_bytes)
                except Exception:
                    logger.error(f"Could not decode or write file from INLINE_DATA, mime_type : {mimeType}")

    mem.seek(0)
    return mem.read()


def send_csv_attachment(filename: str, csv_str: str):
    """
    Sends a single CSV with UTF-8 BOM (so Excel opens it correctly).
    """
    bom_prefixed = "\ufeff" + csv_str
    bio = io.BytesIO(bom_prefixed.encode("utf-8"))
    return send_file(
        bio,
        mimetype="text/csv; charset=utf-8",
        as_attachment=True,
        download_name=filename,
        max_age=0,
    )


def send_zip_attachment(filename: str, zip_bytes: bytes):
    bio = io.BytesIO(zip_bytes)
    return send_file(
        bio,
        mimetype="application/zip",
        as_attachment=True,
        download_name=filename,
        max_age=0,
    )


def _require_user_owned_attachment_path(file_path: str) -> None:
    """Abort unless the current user owns the requested attachment path."""
    current_user = g.get("authIdentifier")
    if not current_user:
        abort(401, description="Authentication required")

    store = get_store()
    if not store.get_derived_document_by_path(file_path):
        logger.warning("Attachment access denied for user %s on path %s", current_user, file_path)
        abort(404, description="Attachment file not found")


def _resolve_admin_asset_folder_id(requested_folder_id: str | None) -> tuple[str, bool]:
    from backend.utils.client_utils import current_request_login
    from backend.utils.permissions_utils import has_project_write_access

    default_id = get_admin_uploads_managedfolder_id()
    login = current_request_login() or ""
    is_admin = has_project_write_access(dataiku.default_project_key(), login)

    if not requested_folder_id:
        if not default_id:
            abort(500, description="Admin uploads managed folder not configured")
        return default_id, is_admin

    if not is_admin and requested_folder_id != default_id:
        logger.warning(f"User {login} attempted to access managed folder {requested_folder_id} without admin rights")
        abort(403, description=f"User {login} does not have permission to access this managed folder")

    return requested_folder_id, is_admin


def _public_admin_asset_paths() -> set[str]:
    from backend.services.admin_config_service import read_config

    cfg = read_config() or {}
    paths = set()
    for key in ("leftPanelLogoPath", "homepageImagePath", "customCssPath"):
        value = cfg.get(key)
        if isinstance(value, dict):
            value = value.get("path")
        if isinstance(value, str) and value:
            paths.add(value)
    return paths


def _build_csv_from_artifacts(artifacts: list[dict], aggregated_folders: dict, source_name):
    total_files = 0
    first_file_candidate = None
    folders = records_folder_from_artifacts(artifacts)
    # { folder: [Table...] }
    for folder_name, tables in folders.items():
        files = build_csv_files_from_records(tables, folder_name)
        if not files:
            continue
        aggregated_folders.setdefault(source_name, {})[folder_name] = files

        logger.info(f"Found {files} CSV files in artifact '{folder_name}'")
        total_files += len(files)
        # Track first CSV file as a candidate for the single-file case
        if first_file_candidate is None:
            # name for single CSV: "<source>-<first_file_name>"
            first_filename = files[0][0]
            csv_str = files[0][1]
            prefix = (source_name or "").strip().replace(" ", "_")
            filename = f"{prefix + '-' if prefix else ''}{first_filename}"
            first_file_candidate = (filename, csv_str)

    return total_files, aggregated_folders, first_file_candidate


def _build_inline_data_file(artifacts: list[schemas.Artifact]):
    """
    extract the inline data files
    also handle SIMPLE_DOCUMENT artf because they come from the same place
    structured agent : generate artifacts
    """
    data_files = []
    simple_doc_counter = 0
    for index, art in enumerate(artifacts) or []:
        for item in art.get("parts") or []:
            part_type = item.get("type")
            if part_type == "DATA_INLINE":
                mimeType = item.get("mimeType", None)
                dataBase64 = item.get("dataBase64", None)
                filename = item.get("filename", None)

                if mimeType and dataBase64:
                    data_files.append(
                        {
                            "mimeType": mimeType,
                            "dataBase64": dataBase64,
                            "filename": filename,
                        }
                    )
                else:
                    logger.debug(f"artifact of type DATA_INLINE could not be extracted : {item}")
            elif part_type == "SIMPLE_DOCUMENT":
                snippet = item.get("markdownSnippet") or item.get("textSnippet")
                if not snippet:
                    logger.debug(f"SIMPLE_DOCUMENT artifact missing snippet: {item}")
                    continue

                if item.get("markdownSnippet"):
                    mimeType = "text/markdown"
                    extension = "md"
                else:
                    mimeType = "text/plain"
                    extension = "txt"

                filename = item.get("filename") or item.get("name")
                if not filename:
                    filename = f"document_{simple_doc_counter}.{extension}"
                simple_doc_counter += 1

                if isinstance(snippet, str):
                    data_files.append(
                        {
                            "mimeType": mimeType,
                            "dataBase64": base64.b64encode(snippet.encode("utf-8")).decode("ascii"),
                            "filename": filename,
                        }
                    )
                else:
                    logger.debug(f"SIMPLE_DOCUMENT artifact snippet is not text: {item}")

    return data_files


def download_artifacts(
    artifacts_meta: dict[str, schemas.ArtifactsMetadata], zip_name: str = "agent-hub-exports.zip"
) -> bytes:
    """
    Returns either:
      - A single CSV (if exactly one table), OR
      - A ZIP of CSVs arranged in folders (if multiple).

    Filename pattern mirrors the frontend:
      - CSV:   "<sourceName><info>-<idx>.csv" (we'll use "<source>-<info>_<idx>.csv")
      - ZIP:   "<source or 'agent-hub-exports'>.zip"
    """

    total_files = 0
    aggregated_csv_folders = {}
    inline_data_files = []

    single_file_candidate = None  # ("csv", filename, csv_str) or ("inline", filename, mime_type, data_base64)
    inline_index = 0
    for meta in (artifacts_meta or {}).values():
        if not isinstance(meta, dict):
            continue
        stored_artifacts = meta.get("artifacts")
        if not isinstance(stored_artifacts, list):
            continue
        artifacts: List[dict] = [
            artifact for artifact in stored_artifacts if isinstance(artifact, dict) and not is_history_marker(artifact)
        ]
        if not artifacts:
            continue
        source_name: str | None = meta.get("agentName")

        total_files_run, aggregated_csv_folders, first_csv_candidate = _build_csv_from_artifacts(
            artifacts, aggregated_csv_folders, source_name
        )
        if single_file_candidate is None and first_csv_candidate:
            filename, csv_str = first_csv_candidate
            single_file_candidate = ("csv", filename, csv_str)
        inline_new_files = _build_inline_data_file(artifacts)
        inline_data_files.extend(inline_new_files)

        total_files += total_files_run
        for item in inline_new_files:
            total_files += 1
            if single_file_candidate is None:
                mime_type = item.get("mimeType")
                data_base64 = item.get("dataBase64")
                filename = item.get("filename")
                if not filename:
                    extension = DIP_MIMETYPE_EXTENSIONS.get(mime_type, "bin")
                    filename = f"document_{inline_index}.{extension}"
                single_file_candidate = ("inline", filename, mime_type, data_base64)
            inline_index += 1

    if total_files == 0:
        abort(404, description="No artifacts found")

    # If exactly 1 file, return it directly for convenience
    if total_files == 1 and single_file_candidate:
        if single_file_candidate[0] == "csv":
            _, filename, csv_str = single_file_candidate
            return send_csv_attachment(filename, csv_str)
        if single_file_candidate[0] == "inline":
            _, filename, mime_type, data_base64 = single_file_candidate
            if mime_type and data_base64:
                try:
                    decoded_bytes = base64.b64decode(data_base64)
                except Exception as e:
                    logger.error(f"Could not decode inline file for download: {e}")
                    abort(500, description="Invalid inline data")
                return send_file(
                    io.BytesIO(decoded_bytes),
                    mimetype=mime_type or "application/octet-stream",
                    as_attachment=True,
                    download_name=filename,
                    max_age=0,
                )

    # Otherwise, zip all CSVs grouped by folder
    zip_bytes = build_zip_bytes(aggregated_csv_folders, inline_data_files)
    return send_zip_attachment(zip_name, zip_bytes)


# ------------------------------------------
# Route: download all artifacts for a message
# ------------------------------------------
@downloads_bp.route("/messages/<msg_id>", methods=["GET"])
@log_http_request
def download_msg_artifacts(msg_id: str):
    # -------- retrieve message  --------------------------------
    store = get_store()
    artifacts_meta = store.get_message_artifacts_meta(msg_id)
    return download_artifacts(artifacts_meta)


@downloads_bp.route("/messages/<msg_id>/artifacts/<art_id>/<art_index>", methods=["GET"])
@log_http_request
def download_artifact_by_id(msg_id: str, art_id: str, art_index: str):
    _, art_grp, art = _get_art(msg_id, art_id, art_index)
    artifacts_meta = {art_id: {**art_grp, "artifacts": [art]}}
    return download_artifacts(artifacts_meta, zip_name=art_grp.get("agentName") or art_id)


@downloads_bp.route("/messages/<msg_id>/artifacts/<art_id>/<art_index>/records/<records_index>", methods=["GET"])
@log_http_request
def download_artifact_records_by_id(msg_id: str, art_id: str, art_index: str, records_index: str):
    art_grp, art, table = _get_rec_tbl(
        msg_id=msg_id,
        art_id=art_id,
        artifact_index=art_index,
        records_index=records_index,
    )
    rec_idx = int(records_index)
    filenames = build_csv_files_from_records(
        [table],
        f"{art.get('name') or 'artifact'}_{rec_idx}",
    )
    filename, csv_str = filenames[0]
    prefix = (art_grp.get("agentName") or "").strip().replace(" ", "_")
    if prefix:
        filename = f"{prefix}-{filename}"
    return send_csv_attachment(filename, csv_str)


@downloads_bp.route("/messages/<msg_id>/artifacts/<art_id>/<art_index>/preview", methods=["GET"])
@log_http_request
def preview_artifact_by_id(msg_id: str, art_id: str, art_index: str):
    try:
        parsed_art_index = int(art_index)
    except ValueError:
        abort(404, description=f"Artifact index '{art_index}' should be a number for artifact id '{art_id}'")

    store = get_store()
    artifacts_meta = store.get_message_artifacts_meta(msg_id)
    try:
        preview = build_generated_document_preview_pdf(artifacts_meta, art_id, parsed_art_index)
    except PreviewArtifactGone as err:
        abort(410, description=str(err))
    except PreviewArtifactNotFound as err:
        abort(404, description=str(err))
    except PreviewConversionError:
        logger.exception(
            "Failed to build generated document preview for message=%s artifact=%s/%s",
            msg_id,
            art_id,
            art_index,
        )
        abort(500, description="Unable to generate preview")

    return send_file(
        io.BytesIO(preview.pdf_bytes),
        mimetype=PDF_MIME_TYPE,
        as_attachment=False,
        download_name=preview.filename,
        max_age=0,
    )


@downloads_bp.route("/attachments", methods=["GET"])
@log_http_request
def download_attachment():
    """
    Stream an attachment directly using its document path.
    Query params: path (required), filename (optional), type (optional)
    """
    from flask import request

    file_path = request.args.get("path")
    if not file_path:
        abort(400, description="Missing 'path' parameter")

    _require_user_owned_attachment_path(file_path)

    filename = request.args.get("filename") or file_path.rsplit("/", 1)[-1]
    mimetype = request.args.get("type") or "application/octet-stream"

    folder_id = get_uploads_managedfolder_id()
    try:
        folder = dataiku.Folder(folder_id)
    except Exception as e:
        logger.exception("Unable to open managed folder %s: %s", folder_id, e)
        abort(500, description="Unable to access attachment storage")

    try:
        with folder.get_download_stream(file_path) as stream:
            data = stream.read()
    except Exception as e:
        logger.exception("Failed to read attachment from path %s: %s", file_path, e)
        abort(404, description="Attachment file not found")

    return send_file(
        io.BytesIO(data),
        as_attachment=True,
        download_name=filename,
        mimetype=mimetype,
        max_age=0,
    )


@downloads_bp.route("/admin/assets", methods=["GET"])
@log_http_request
def download_admin_asset():
    """
    Stream an admin asset (logos, images, etc.) directly using its document path.
    Query params: path (required), managedFolderId (optional), filename (optional), type (optional)

    If managedFolderId is provided, uses that folder only if user is admin(for draft/preview scenarios).
    Otherwise falls back to the persisted admin uploads managed folder from database.
    """
    from flask import request

    file_path = request.args.get("path")
    if not file_path:
        abort(400, description="Missing 'path' parameter")

    filename = request.args.get("filename") or file_path.rsplit("/", 1)[-1]

    # Auto-detect MIME type from file extension, or use provided type
    mimetype = request.args.get("type")
    if not mimetype:
        # Guess MIME type from filename
        guessed_type, _ = mimetypes.guess_type(filename)
        mimetype = guessed_type or "application/octet-stream"
    folder_id, is_admin = _resolve_admin_asset_folder_id(request.args.get("managedFolderId"))

    if not is_admin and file_path not in _public_admin_asset_paths():
        abort(403, description="Asset path is not publicly accessible")

    try:
        folder = dataiku.Folder(folder_id)
    except Exception as e:
        logger.exception("Unable to open admin uploads managed folder %s: %s", folder_id, e)
        abort(500, description="Unable to access admin asset storage")

    try:
        with folder.get_download_stream(file_path) as stream:
            data = stream.read()
    except Exception as e:
        logger.exception("Failed to read admin asset from path %s: %s", file_path, e)
        abort(404, description="Admin asset file not found")

    return send_file(
        io.BytesIO(data),
        as_attachment=False,
        download_name=filename,
        mimetype=mimetype,
    )
