from __future__ import annotations

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

import dataiku
from flask import Blueprint, abort, send_file

from backend.models.artifacts import Artifact, ArtifactsMetadata
from backend.utils.logger_utils import log_http_request
from backend.utils.logging_utils import get_logger
from backend.utils.utils import get_store
from backend.config import get_uploads_managedfolder_id

logger = get_logger(__name__)

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


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


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(table.columns)
    for row in table.data:
        writer.writerow(row)
    return output.getvalue()


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


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 _safe(s: str) -> str:
    return (s or "").strip().replace(" ", "_").replace("/", "_")


def build_zip_bytes(aggregated_folders: Dict[str, Dict[str, List]]) -> bytes:
    """
    folders: { folderName: [(filename, csv_str), ...], ... }
    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)

    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 download_artifacts(artifacts_meta: dict[str, 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"
    """
    aggregated_folders = {}
    single_file_candidate: Tuple[str, str] | None = None  # (filename, csv_str)
    total_files = 0
    for meta in (artifacts_meta or {}).values():
        artifacts: List[dict] = meta.get("artifacts") or []
        if not artifacts:
            continue
        source_name: str | None = meta.get("agentName")
        # if source_name and source_name.strip():
        #     zip_name = f"{source_name}.zip"
        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 single-file case
            if total_files == 1:
                # 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}"
                single_file_candidate = (filename, csv_str)

    if total_files == 0:
        abort(404, description="No artifacts found")
    # If exactly 1 CSV, return it directly for convenience
    if total_files == 1 and single_file_candidate:
        filename, csv_str = single_file_candidate
        return send_csv_attachment(filename, csv_str)

    # Otherwise, zip all CSVs grouped by folder
    zip_bytes = build_zip_bytes(aggregated_folders)
    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):
    store = get_store()
    artifacts_meta = store.get_message_artifacts_meta(msg_id)
    if not artifacts_meta or art_id not in artifacts_meta:
        abort(404, description=f"Artifact id '{art_id}' not found for message '{msg_id}'")
    artifacts = artifacts_meta[art_id].get("artifacts") or []
    if int(art_index) < 0 or int(art_index) >= len(artifacts):
        abort(404, description=f"Artifact index '{art_index}' out of range for artifact id '{art_id}'")
    # Keep only the selected artifact
    artifacts_meta = {art_id: {**artifacts_meta[art_id], "artifacts": [artifacts[int(art_index)]]}}
    return download_artifacts(artifacts_meta, zip_name=artifacts_meta[art_id]["agentName"])


@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")
    
    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,
    )
