#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
generate_all_batches_from_dict_v1.py

Generates blockchain batch JSON + QR files for all lineages:
P. vannamei, P. monodon, and P. indicus.

Author: CIBA IoT Traceability Project
Version: 2025.10
"""

import os
import re
import json
import random
import qrcode
from datetime import datetime, timezone
from urllib.parse import quote_plus
from blockchain_logic import (
    LINEAGE_P_VANNAMEI,
    LINEAGE_P_MONODON,
    LINEAGE_P_INDICUS
)

# ---------------- Configuration ----------------
from config import PUBLIC_BASE_URL

BASE_STAGE_URL = f"{PUBLIC_BASE_URL}/stage"
BASE_PRODUCT_URL = f"{PUBLIC_BASE_URL}/product"

OUT_DATA_DIR = os.path.join("static", "data")
OUT_QR_DIR = os.path.join("static", "qr")

os.makedirs(OUT_DATA_DIR, exist_ok=True)
os.makedirs(OUT_QR_DIR, exist_ok=True)

# ---------------- Upload Stage Map ----------------
UPLOAD_ALLOWED_STAGES = [
    "Broodstock Facility",
    "Hatchery/Nursery",
    "Processing (direct from farms)-1",
    "Processing (direct from farms)-2",
    "Exporting-1",
    "Exporting-2",
    "Retailing-1",
    "Retailing-2",
    "Retailing-3"
]

# ---------------- Canonical Stage List ----------------
CTE_STAGES = [
    "Broodstock Facility",
    "Hatchery/Nursery",
    "Grow-out System",
    "Production of feed & inputs",
    "Harvesting",
    "Aggregation/Trading at Farm Gate",
    "Transportation",
    "Marketing",
    "Processing (direct from farms)",
    "Exporting",
    "Retailing",
    "Shrimp on a plate",
    "Consumption / Feedback"
]


# ---------------- Utility Functions ----------------
def normalize_stage_name(name: str) -> str:
    """Normalize stage names for comparison."""
    return (
        name.lower()
        .replace("&", "and")
        .replace("/", " ")
        .replace("-", " ")
        .replace("(", "")
        .replace(")", "")
        .strip()
    )


def sanitize_stage_label(s: str) -> str:
    """Convert to safe filename."""
    return re.sub(r"[^A-Za-z0-9_-]", "-", s.strip()).replace("--", "-")


def create_qr_file(url: str, file_path: str):
    """Generate QR code for a given URL."""
    os.makedirs(os.path.dirname(file_path), exist_ok=True)
    qr = qrcode.QRCode(
        error_correction=qrcode.constants.ERROR_CORRECT_L,
        box_size=6,
        border=2
    )
    qr.add_data(url)
    qr.make(fit=True)
    qr.make_image(fill_color="black", back_color="white").save(file_path)
    return file_path


def build_stage_map(lineage_list):
    """Create normalized lookup dict."""
    stage_map = {}
    for entry in lineage_list:
        stage_name = entry.get("stage", "").strip()
        if not stage_name:
            continue
        key = normalize_stage_name(stage_name)
        stage_map.setdefault(key, []).append(entry)
    return stage_map


# ---------------- Batch Builder ----------------
def build_batches(species_name, lineage_data, n_batches=10):
    """Builds batch files for a specific lineage."""
    species_dir_data = os.path.join(OUT_DATA_DIR, species_name.lower())
    species_dir_qr = os.path.join(OUT_QR_DIR, species_name.lower())
    os.makedirs(species_dir_data, exist_ok=True)
    os.makedirs(species_dir_qr, exist_ok=True)

    stage_map = build_stage_map(lineage_data)

    for i in range(1, n_batches + 1):
        batch_id = f"{species_name[:3].upper()}-CIBA-{i:03d}"
        out_json = os.path.join(species_dir_data, f"{batch_id}.json")

        lineage = []
        print(f"\n?? Building {batch_id} ({species_name})")

        for stage in CTE_STAGES[:-1]:
            norm_key = normalize_stage_name(stage)
            matched_keys = [k for k in stage_map if norm_key in k or k in norm_key]

            if matched_keys:
                chosen_key = random.choice(matched_keys)
                data_entry = random.choice(stage_map[chosen_key])
                tx = data_entry
                stage_display = data_entry.get("stage", stage)
                print(f"? Matched: {stage} -> {stage_display}")
            else:
                tx = {"Auto-Filled": "Yes", "Note": f"No data found for {stage}"}
                stage_display = stage
                print(f"?? Missing data, filled placeholder: {stage}")

            qr_safe_stage = quote_plus(stage_display)
            qr_name = f"{batch_id}_{sanitize_stage_label(stage_display)}.png"
            qr_rel = f"qr/{species_name.lower()}/{qr_name}"
            qr_path = os.path.join(species_dir_qr, qr_name)
            qr_link = f"{BASE_STAGE_URL}/{batch_id}/{qr_safe_stage}"
            create_qr_file(qr_link, qr_path)

            upload_allowed = stage_display in UPLOAD_ALLOWED_STAGES

            lineage.append({
                "Stage": stage_display,
                "Timestamp": datetime.now(timezone.utc).isoformat(),
                "TransactionData": tx,
                "QR": qr_rel,
                "QR_Link": qr_link,
                "CertificateUpload": upload_allowed
            })

        # Feedback stage (reuse main QR)
        fb_stage = "Consumption / Feedback"
        main_qr_link = f"{BASE_PRODUCT_URL}/{batch_id}"
        main_qr_name = f"{batch_id}_MAIN.png"
        main_qr_rel = f"qr/{species_name.lower()}/{main_qr_name}"
        create_qr_file(main_qr_link, os.path.join(species_dir_qr, main_qr_name))
        lineage.append({
            "Stage": fb_stage,
            "Timestamp": datetime.now(timezone.utc).isoformat(),
            "TransactionData": {"Note": "Consumer feedback placeholder"},
            "QR": main_qr_rel,
            "QR_Link": main_qr_link,
            "CertificateUpload": False
        })

        batch_obj = {
            "product_id": batch_id,
            "species": f"P. {species_name}",
            "generated_on": datetime.now(timezone.utc).isoformat(),
            "lineage": lineage
        }

        with open(out_json, "w", encoding="utf-8") as f:
            json.dump(batch_obj, f, indent=2, ensure_ascii=False)

        print(f"?? Saved {batch_id} with {len(lineage)} stages ? {out_json}")

    print(f"\n? All {species_name} batches built successfully!\n")


# ---------------- Main ----------------
def main():
    lineage_sets = {
        "vannamei": LINEAGE_P_VANNAMEI,
        "monodon": LINEAGE_P_MONODON,
        "indicus": LINEAGE_P_INDICUS
    }

    print("?? Starting batch generation for all lineages...")

    for species, lineage_data in lineage_sets.items():
        print(f"\n?? Processing lineage: {species.capitalize()}")
        build_batches(species, lineage_data, n_batches=10)

    print("\n?? All lineage batches generated successfully!")


if __name__ == "__main__":
    main()
