#!/usr/bin/env julia
#############################################################################
## 06_save_model.jl
## CSV-First Save Architecture
## Robust, portable, no serialization issues
#############################################################################

module RobustSave

using Dates, Printf, CSV, DataFrames, JSON

export robust_save_model_csv, robust_save_model

"""
CSV-First Save Architecture

When chains_already_saved=true (BULLETPROOF MODE):
- Chains already saved to run_dir/chains/ by model execution
- Just verify chains and add metadata/data files
- Even if this function crashes, chains are SAFE

When chains_already_saved=false (legacy mode):
- Copy chains from temp directory to new run directory
- Add metadata/data files

Saves model results as:
1. CSV chain files (source of truth - never fail)
2. Data CSVs (original inputs for reproducibility)
3. Simple metadata.json (no complex types)
4. Human-readable README.txt

No JLD2, no serialization issues, fully portable and reproducible.
"""
function robust_save_model_csv(
    run_dir_or_temp::String,
    data_dict::Dict,
    original_data::Dict,
    metadata::Dict;
    chains_already_saved::Bool=false
)
    println("\n" * "=" ^ 70)
    if chains_already_saved
        println("ADDING METADATA TO EXISTING RUN (chains already secured)")
    else
        println("CSV-FIRST MODEL SAVE")
    end
    println("=" ^ 70)

    # Determine directories based on mode
    if chains_already_saved
        # Chains already saved - run_dir_or_temp IS the run directory
        run_dir = run_dir_or_temp
        chains_dir = joinpath(run_dir, "chains")
        data_dir = joinpath(run_dir, "data")
        run_id = basename(run_dir)
        timestamp = replace(run_id, "run_" => "")

        println("Run directory: $run_dir")
        println("Mode: Chains already secured, adding metadata")

        # Verify chains directory exists
        if !isdir(chains_dir)
            error("CRITICAL: Chains directory not found: $chains_dir")
        end

        # Count existing chain files
        chain_files = filter(f -> endswith(f, ".csv") && contains(f, "chain"), readdir(chains_dir))
        if isempty(chain_files)
            error("CRITICAL: No chain CSV files found in $chains_dir")
        end

        println("Found $(length(chain_files)) chain files already saved")

        # Calculate total size
        total_size_gb = 0.0
        for chain_file in chain_files
            chain_path = joinpath(chains_dir, chain_file)
            total_size_gb += filesize(chain_path) / (1024^3)
        end

        println("✓ Chains verified ($(round(total_size_gb, digits=2)) GB total)")

    else
        # Legacy mode - create new run directory and copy chains
        temp_csv_dir = run_dir_or_temp

        timestamp = Dates.format(Dates.now(), "yyyy-mm-dd_HH-MM-SS")
        run_id = "run_$(timestamp)"
        run_dir = joinpath("outputs", "model_outputs", "latest", run_id)
        chains_dir = joinpath(run_dir, "chains")
        data_dir = joinpath(run_dir, "data")

        println("Run ID: $run_id")
        println("Output directory: $run_dir")

        # Create directory structure
        mkpath(chains_dir)

        # STEP 1: Copy CSV chain files
        println("\n" * "=" ^ 70)
        println("STEP 1: Copying MCMC chain CSV files")
        println("=" ^ 70)
        println("Source: $temp_csv_dir")
        println("Destination: $chains_dir")

        csv_files = filter(f -> endswith(f, ".csv"), readdir(temp_csv_dir))
        chain_files = filter(f -> contains(f, "chain"), csv_files)

        if isempty(chain_files)
            error("No chain CSV files found in $temp_csv_dir")
        end

        println("Found $(length(chain_files)) chain files")

        total_size_gb = 0.0
        for (i, csv_file) in enumerate(sort(chain_files))
            src_path = joinpath(temp_csv_dir, csv_file)

            # Rename to standard format: chain_1.csv, chain_2.csv, etc.
            dest_filename = "chain_$i.csv"
            dest_path = joinpath(chains_dir, dest_filename)

            src_size = filesize(src_path)
            size_gb = src_size / (1024^3)
            total_size_gb += size_gb

            println("  Copying $csv_file → $dest_filename ($(round(size_gb, digits=2)) GB)")
            cp(src_path, dest_path, force=true)

            # Verify copy with size check
            dest_size = filesize(dest_path)
            if dest_size != src_size
                error("CRITICAL: Size mismatch for $dest_filename! Source: $src_size, Dest: $dest_size")
            end
        end

        println("✓ All chains copied and verified ($(round(total_size_gb, digits=2)) GB total)")
    end

    # Create data directory
    mkpath(data_dir)

    # STEP 2: Save data CSVs
    println("\n" * "=" ^ 70)
    println("STEP 2: Saving original data CSVs")
    println("=" ^ 70)

    for (name, df) in original_data
        if isa(df, DataFrame)
            csv_path = joinpath(data_dir, "$(name).csv")
            println("  Saving $(name).csv ($(nrow(df)) rows)")
            CSV.write(csv_path, df)
        end
    end

    println("✓ Data CSVs saved")

    # STEP 3: Save Stan data dictionary as JSON
    println("\n" * "=" ^ 70)
    println("STEP 3: Saving Stan data dictionary")
    println("=" ^ 70)

    # Convert data_dict to JSON-serializable format
    stan_data_json = Dict{String, Any}()
    for (k, v) in data_dict
        try
            # Only save simple types (numbers, arrays of numbers)
            if isa(v, Number) || isa(v, AbstractArray{<:Number})
                stan_data_json[k] = v
            elseif isa(v, AbstractArray)
                # Try to convert, skip if fails
                try
                    stan_data_json[k] = collect(v)
                catch
                    println("  Skipping $k (complex type)")
                end
            end
        catch e
            println("  Warning: Could not serialize $k: $e")
        end
    end

    stan_data_path = joinpath(data_dir, "stan_data.json")
    open(stan_data_path, "w") do f
        JSON.print(f, stan_data_json, 2)
    end
    println("✓ Stan data saved to stan_data.json")

    # Count chain files for metadata
    chain_files_final = filter(f -> endswith(f, ".csv") && contains(f, "chain"), readdir(chains_dir))
    num_chains = length(chain_files_final)

    # Recalculate total_size_gb if in chains_already_saved mode
    if chains_already_saved
        total_size_gb = 0.0
        for chain_file in chain_files_final
            chain_path = joinpath(chains_dir, chain_file)
            total_size_gb += filesize(chain_path) / (1024^3)
        end
    end

    # STEP 4: Save metadata
    println("\n" * "=" ^ 70)
    println("STEP 4: Saving metadata")
    println("=" ^ 70)

    # Add run info to metadata
    metadata["run_id"] = run_id
    metadata["timestamp"] = timestamp
    metadata["files"] = Dict(
        "chains" => ["chains/chain_$i.csv" for i in 1:num_chains],
        "data" => readdir(data_dir),
        "chain_size_gb" => num_chains > 0 ? round(total_size_gb / num_chains, digits=2) : 0.0,
        "total_size_gb" => round(total_size_gb, digits=2)
    )

    metadata_path = joinpath(run_dir, "metadata.json")
    open(metadata_path, "w") do f
        JSON.print(f, metadata, 2)
    end
    println("✓ Metadata saved to metadata.json")

    # STEP 5: Generate README
    println("\n" * "=" ^ 70)
    println("STEP 5: Generating README")
    println("=" ^ 70)

    readme_path = joinpath(run_dir, "README.txt")
    generate_readme(readme_path, run_id, metadata, num_chains, total_size_gb)
    println("✓ README generated")

    # STEP 6: Final verification
    println("\n" * "=" ^ 70)
    println("STEP 6: Verification")
    println("=" ^ 70)

    # Verify all chain files exist and are readable
    all_good = true
    verified_files = Dict{String, Dict{String, Any}}()

    for i in 1:num_chains
        chain_path = joinpath(chains_dir, "chain_$i.csv")
        if !isfile(chain_path)
            println("  ✗ Missing: chain_$i.csv")
            all_good = false
        else
            # Quick read test and size check
            try
                CSV.File(chain_path; limit=1)
                file_size_gb = filesize(chain_path) / (1024^3)
                verified_files["chain_$i.csv"] = Dict(
                    "path" => chain_path,
                    "size_gb" => file_size_gb,
                    "verified" => true
                )
                println("  ✓ chain_$i.csv verified ($(round(file_size_gb, digits=2)) GB)")
            catch e
                println("  ✗ Cannot read chain_$i.csv: $e")
                all_good = false
            end
        end
    end

    if !all_good
        error("Verification failed - some files are missing or corrupted")
    end

    println("\n" * "=" ^ 70)
    println("✓ MODEL SAVED SUCCESSFULLY")
    println("=" ^ 70)
    println("Run directory: $run_dir")
    println("Total size: $(round(total_size_gb, digits=2)) GB")
    println("Status: All files verified and ready")
    println("=" ^ 70)

    # Return verification details for cleanup
    return (
        run_dir = run_dir,
        verified_files = verified_files,
        total_size_gb = total_size_gb,
        verification_passed = all_good,
        num_chains = num_chains
    )
end

function generate_readme(
    filepath::String,
    run_id::String,
    metadata::Dict,
    num_chains::Int,
    total_size_gb::Float64
)
    """Generate human-readable README file"""

    open(filepath, "w") do f
        write(f, "=" ^ 78 * "\n")
        write(f, "4D LATENT TRAIT MODEL - MODEL RUN RESULTS\n")
        write(f, "=" ^ 78 * "\n\n")

        write(f, "Run ID: $run_id\n")
        write(f, "Model: $(get(metadata, "model_file", "unknown"))\n")
        write(f, "Date: $(Dates.format(Dates.now(), "yyyy-mm-dd HH:MM:SS"))\n")
        write(f, "Status: $(get(metadata, "convergence_status", "unknown"))\n\n")

        write(f, "=" ^ 78 * "\n")
        write(f, "DIRECTORY CONTENTS\n")
        write(f, "=" ^ 78 * "\n\n")

        write(f, "chains/\n")
        for i in 1:num_chains
            write(f, "  ├── chain_$i.csv\n")
        end
        write(f, "  Total: $(get(metadata, "num_chains", num_chains)) chains × " *
              "$(get(metadata, "num_samples", "?")) samples\n")
        write(f, "  Size: $(round(total_size_gb, digits=2)) GB\n\n")

        write(f, "data/\n")
        write(f, "  ├── text_data.csv\n")
        write(f, "  ├── expert_dim.csv\n")
        write(f, "  ├── expert_lr.csv\n")
        write(f, "  ├── segment_year_map.csv (V10)\n")
        write(f, "  ├── segment_info.csv (V10)\n")
        write(f, "  └── stan_data.json\n\n")

        write(f, "=" ^ 78 * "\n")
        write(f, "MODEL CONFIGURATION\n")
        write(f, "=" ^ 78 * "\n\n")

        write(f, "Chains: $(get(metadata, "num_chains", "?"))\n")
        write(f, "Warmup: $(get(metadata, "num_warmup", "?"))\n")
        write(f, "Samples: $(get(metadata, "num_samples", "?"))\n")
        write(f, "Adapt delta: $(get(metadata, "adapt_delta", "?"))\n")
        write(f, "Max depth: $(get(metadata, "max_depth", "?"))\n\n")

        write(f, "Dimensions: $(join(get(metadata, "dimensions", ["?"]), ", "))\n\n")

        write(f, "=" ^ 78 * "\n")
        write(f, "HOW TO USE THESE RESULTS\n")
        write(f, "=" ^ 78 * "\n\n")

        write(f, "To extract party positions:\n\n")
        write(f, "  julia 02_post_estimation.jl\n\n")

        write(f, "This will read the CSV files and generate party_positions_[timestamp].csv\n")
        write(f, "with uncertainty estimates (SE, credible intervals).\n\n")

        write(f, "=" ^ 78 * "\n")
        write(f, "Generated: $(Dates.format(Dates.now(), "yyyy-mm-dd HH:MM:SS"))\n")
        write(f, "=" ^ 78 * "\n")
    end
end

"""
Wrapper for robust_save_model_csv that handles the stanmodel tuple from run_4dim_stan_model.

The model execution now saves chains BEFORE returning, so this function:
1. Verifies chains are already saved in final_run_dir
2. Adds metadata and data files
3. Returns the output path

Arguments:
- stanmodel_tuple: Named tuple from run_4dim_stan_model (contains .stanmodel, .final_run_dir, etc.)
- model_data: Dict with data_dict, original data, and metadata
- output_dir: Base output directory (ignored - uses stanmodel_tuple.final_run_dir)
- compress: Ignored (CSV-first architecture)
- keep_local_backups: Ignored (chains already saved)
"""
function robust_save_model(
    stanmodel_tuple,
    model_data::Dict,
    output_dir::String;
    compress::Bool=true,
    keep_local_backups::Int=2
)
    # Extract the final run directory from the stanmodel tuple
    if !hasproperty(stanmodel_tuple, :final_run_dir)
        error("stanmodel_tuple missing :final_run_dir - chains may not be saved!")
    end

    final_run_dir = stanmodel_tuple.final_run_dir

    # Prepare original data for saving
    original_data = Dict{String, Any}()
    if haskey(model_data, "manifesto")
        original_data["text_data"] = model_data["manifesto"]
    end
    if haskey(model_data, "expert_dim")
        original_data["expert_dim"] = model_data["expert_dim"]
    end
    if haskey(model_data, "expert_lr")
        original_data["expert_lr"] = model_data["expert_lr"]
    end
    # V10: Save segment_year_map and segment_info for post-estimation
    if haskey(model_data, "segment_year")
        original_data["segment_year_map"] = model_data["segment_year"]
    end
    if haskey(model_data, "segment_info")
        original_data["segment_info"] = model_data["segment_info"]
    end
    # V9 fallback: Save party_year_map for post-estimation (includes interpolated years)
    if haskey(model_data, "party_year")
        original_data["party_year_map"] = model_data["party_year"]
    end

    # Prepare metadata
    metadata = Dict{String, Any}()
    if haskey(model_data, "model_info")
        for (k, v) in model_data["model_info"]
            metadata[string(k)] = v
        end
    end

    # Get data_dict
    data_dict = get(model_data, "data_dict", Dict{String, Any}())

    # Call the CSV-first save with chains_already_saved=true
    result = robust_save_model_csv(
        final_run_dir,
        data_dict,
        original_data,
        metadata;
        chains_already_saved=true
    )

    return result.run_dir
end

end  # module
