#include "args.h"
#include "bromberg_sl2.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

typedef struct {
    hash_matrix_t expected_doc;
    hash_matrix_t prefix;
    hash_matrix_t quote;
    hash_matrix_t suffix;
} SequinclusionProof;

// Get indexing mode (0 or 1 indexed)
bool get_indexing_mode(int index_base) {
    // 1. Priority: CLI Flag
    if (index_base != -1) {
        return index_base == 1;
    }

    // 2. Priority: Environment Variable
    char* env_val = getenv("POTC_INDEX_BASE");
    if (env_val) {
        int val = atoi(env_val);
        if (val == 1) return true;
        if (val == 0) return false;
        fprintf(stderr, "error: invalid value in $POTC_INDEX_BASE: '%s' (must be 0 or 1)\n", env_val);
        exit(1);
    }

    // 3. Priority: Default
    return false;  // 0-indexed
}

// Parse hash matrix from hex string (64 hex chars)
int hash_matrix_from_hex(const char* hex, hash_matrix_t* out) {
    if (strlen(hex) < 64) {
        fprintf(stderr, "error: expected at least 64 hex chars for hash matrix\n");
        return -1;
    }

    for (int i = 0; i < 4; i++) {
        char chunk[17];
        strncpy(chunk, hex + i * 16, 16);
        chunk[16] = '\0';
        out->m[i] = strtoull(chunk, NULL, 16);
    }

    return 0;
}

// Parse sequinclusion proof from hex string (256 hex chars)
int parse_sequinclusion_proof(const char* input, SequinclusionProof* proof) {
    size_t len = strlen(input);
    if (len < 256) {
        fprintf(stderr, "error: expected at least 256 hex chars for sequinclusion proof, got %zu\n", len);
        return -1;
    }

    if (hash_matrix_from_hex(input + 0, &proof->expected_doc) != 0) return -1;
    if (hash_matrix_from_hex(input + 64, &proof->prefix) != 0) return -1;
    if (hash_matrix_from_hex(input + 128, &proof->quote) != 0) return -1;
    if (hash_matrix_from_hex(input + 192, &proof->suffix) != 0) return -1;

    return 0;
}

// Verify sequinclusion proof
bool verify_sequinclusion(const SequinclusionProof* proof) {
    hash_matrix_t computed = bromberg_compose(
        bromberg_compose(proof->prefix, proof->quote),
        proof->suffix
    );
    return hash_matrix_eq(&computed, &proof->expected_doc);
}

// Compute and format proof
int compute_and_format_proof(const uint8_t* contents, size_t len, size_t start, size_t end, bool one_indexed, char* out) {
    // Convert to 0-indexed internally
    size_t start_0, end_exclusive;
    if (one_indexed) {
        if (start == 0) {
            fprintf(stderr, "error: start index must be >= 1 in 1-indexed mode\n");
            return -1;
        }
        start_0 = start - 1;
        end_exclusive = end;
    } else {
        start_0 = start;
        end_exclusive = end + 1;
    }

    if (end_exclusive > len) {
        fprintf(stderr, "error: end index exceeds content length %zu\n", len);
        return -1;
    }

    hash_matrix_t doc_hash = bromberg_hash(contents, len);
    hash_matrix_t prefix_hash = (start_0 == 0) ? MATRIX_I : bromberg_hash(contents, start_0);
    hash_matrix_t quote_hash = bromberg_hash(contents + start_0, end_exclusive - start_0);
    hash_matrix_t suffix_hash = (end_exclusive == len) ? MATRIX_I : bromberg_hash(contents + end_exclusive, len - end_exclusive);

    char hex_doc[65], hex_prefix[65], hex_quote[65], hex_suffix[65];
    bromberg_to_hex(doc_hash, hex_doc);
    bromberg_to_hex(prefix_hash, hex_prefix);
    bromberg_to_hex(quote_hash, hex_quote);
    bromberg_to_hex(suffix_hash, hex_suffix);

    sprintf(out, "%s%s%s%s", hex_doc, hex_prefix, hex_quote, hex_suffix);
    return 0;
}

// Compute sequinclusion proof struct
int compute_sequinclusion_proof(const uint8_t* contents, size_t len, size_t start, size_t end, bool one_indexed, SequinclusionProof* proof) {
    size_t start_0, end_exclusive;
    if (one_indexed) {
        if (start == 0) {
            fprintf(stderr, "error: start index must be >= 1 in 1-indexed mode\n");
            return -1;
        }
        start_0 = start - 1;
        end_exclusive = end;
    } else {
        start_0 = start;
        end_exclusive = end + 1;
    }

    if (end_exclusive > len) {
        fprintf(stderr, "error: end index exceeds content length %zu\n", len);
        return -1;
    }

    proof->expected_doc = bromberg_hash(contents, len);
    proof->prefix = (start_0 == 0) ? MATRIX_I : bromberg_hash(contents, start_0);
    proof->quote = bromberg_hash(contents + start_0, end_exclusive - start_0);
    proof->suffix = (end_exclusive == len) ? MATRIX_I : bromberg_hash(contents + end_exclusive, len - end_exclusive);

    return 0;
}

// Read file contents
uint8_t* read_file(const char* path, size_t* len) {
    FILE* f = fopen(path, "rb");
    if (!f) {
        fprintf(stderr, "error: failed to read file '%s'\n", path);
        return NULL;
    }

    fseek(f, 0, SEEK_END);
    *len = ftell(f);
    fseek(f, 0, SEEK_SET);

    uint8_t* data = malloc(*len);
    if (!data) {
        fprintf(stderr, "error: out of memory\n");
        fclose(f);
        return NULL;
    }

    fread(data, 1, *len, f);
    fclose(f);
    return data;
}

void format_proof(const SequinclusionProof* proof, char* out) {
    char hex_doc[65], hex_prefix[65], hex_quote[65], hex_suffix[65];
    bromberg_to_hex(proof->expected_doc, hex_doc);
    bromberg_to_hex(proof->prefix, hex_prefix);
    bromberg_to_hex(proof->quote, hex_quote);
    bromberg_to_hex(proof->suffix, hex_suffix);
    sprintf(out, "%s%s%s%s", hex_doc, hex_prefix, hex_quote, hex_suffix);
}

void do_generate_sequinclusion(Args* args) {
    bool one_indexed = get_indexing_mode(args->index_base);

    if (args->index_count != 1) {
        fprintf(stderr, "error: sequinclusion generate requires exactly one -i flag\n");
        exit(1);
    }

    uint8_t* contents;
    size_t len;

    if (args->file_count == 1) {
        contents = read_file(args->files[0], &len);
        if (!contents) exit(1);
    } else if (args->text_count == 1) {
        contents = (uint8_t*)args->texts[0];
        len = strlen(args->texts[0]);
    } else {
        fprintf(stderr, "error: sequinclusion generate requires exactly one -f or -t\n");
        exit(1);
    }

    char proof[257];
    if (compute_and_format_proof(contents, len, args->indices[0].start, args->indices[0].end, one_indexed, proof) != 0) {
        if (args->file_count == 1) free(contents);
        exit(1);
    }

    printf("%s\n", proof);
    if (args->file_count == 1) free(contents);
}

void do_verify_sequinclusion(Args* args) {
    for (int i = 0; i < args->proof_count; i++) {
        SequinclusionProof proof;
        if (parse_sequinclusion_proof(args->proofs[i], &proof) != 0) {
            if (args->proof_count == 1) {
                printf("failed sequinclusion\n");
            } else {
                printf("failed sequinclusion %d\n", i + 1);
            }
            continue;
        }

        if (verify_sequinclusion(&proof)) {
            if (args->proof_count == 1) {
                printf("verified sequinclusion\n");
            } else {
                printf("verified sequinclusion %d\n", i + 1);
            }
        } else {
            if (args->proof_count == 1) {
                printf("failed sequinclusion\n");
            } else {
                printf("failed sequinclusion %d\n", i + 1);
            }
        }
    }
}

void do_generate_transclusion(Args* args) {
    bool one_indexed = get_indexing_mode(args->index_base);

    int total_sources = args->file_count + args->text_count;
    if (total_sources != 2) {
        fprintf(stderr, "error: transclusion generate requires exactly 2 sources (got %d)\n", total_sources);
        exit(1);
    }

    if (args->index_count != 2) {
        fprintf(stderr, "error: number of sources (-f/-t) must match number of indices (-i). got %d sources, %d indices\n", total_sources, args->index_count);
        exit(1);
    }

    uint8_t* contents1;
    uint8_t* contents2;
    size_t len1, len2;
    bool free1 = false, free2 = false;

    // Get first source
    if (args->file_count > 0) {
        contents1 = read_file(args->files[0], &len1);
        if (!contents1) exit(1);
        free1 = true;
    } else {
        contents1 = (uint8_t*)args->texts[0];
        len1 = strlen(args->texts[0]);
    }

    // Get second source
    if (args->file_count == 2) {
        contents2 = read_file(args->files[1], &len2);
        if (!contents2) {
            if (free1) free(contents1);
            exit(1);
        }
        free2 = true;
    } else if (args->file_count == 1 && args->text_count == 1) {
        contents2 = (uint8_t*)args->texts[0];
        len2 = strlen(args->texts[0]);
    } else {
        contents2 = (uint8_t*)args->texts[1];
        len2 = strlen(args->texts[1]);
    }

    SequinclusionProof proof1, proof2;
    if (compute_sequinclusion_proof(contents1, len1, args->indices[0].start, args->indices[0].end, one_indexed, &proof1) != 0) {
        if (free1) free(contents1);
        if (free2) free(contents2);
        exit(1);
    }

    if (compute_sequinclusion_proof(contents2, len2, args->indices[1].start, args->indices[1].end, one_indexed, &proof2) != 0) {
        if (free1) free(contents1);
        if (free2) free(contents2);
        exit(1);
    }

    bool inc1_ok = verify_sequinclusion(&proof1);
    bool inc2_ok = verify_sequinclusion(&proof2);
    bool quotes_match = hash_matrix_eq(&proof1.quote, &proof2.quote);

    if (inc1_ok && inc2_ok && quotes_match) {
        char proof_str1[257], proof_str2[257];
        format_proof(&proof1, proof_str1);
        format_proof(&proof2, proof_str2);
        printf("%s%s\n", proof_str1, proof_str2);
    } else {
        fprintf(stderr, "error: invalid transclusion. reason: ");
        if (!inc1_ok) fprintf(stderr, "(sequinclusion 1 failed)");
        if (!inc2_ok) fprintf(stderr, "(sequinclusion 2 failed)");
        if (!quotes_match) fprintf(stderr, "(transclusion failed (quotes don't match))");
        fprintf(stderr, "\n");
        if (free1) free(contents1);
        if (free2) free(contents2);
        exit(1);
    }

    if (free1) free(contents1);
    if (free2) free(contents2);
}

void do_verify_transclusion_proofs(Args* args) {
    for (int i = 0; i < args->proof_count; i++) {
        const char* input = args->proofs[i];
        size_t len = strlen(input);

        if (len != 512) {
            fprintf(stderr, "error: transclusion %d expected 512 hex chars, got %zu\n", i + 1, len);
            if (args->proof_count == 1) {
                printf("failed transclusion. reason: (invalid input length)\n");
            } else {
                printf("failed transclusion %d. reason: (invalid input length)\n", i + 1);
            }
            continue;
        }

        SequinclusionProof proof1, proof2;
        if (parse_sequinclusion_proof(input, &proof1) != 0) {
            if (args->proof_count == 1) {
                printf("failed transclusion. reason: (parse error proof 1)\n");
            } else {
                printf("failed transclusion %d. reason: (parse error proof 1)\n", i + 1);
            }
            continue;
        }

        if (parse_sequinclusion_proof(input + 256, &proof2) != 0) {
            if (args->proof_count == 1) {
                printf("failed transclusion. reason: (parse error proof 2)\n");
            } else {
                printf("failed transclusion %d. reason: (parse error proof 2)\n", i + 1);
            }
            continue;
        }

        bool inc1_ok = verify_sequinclusion(&proof1);
        bool inc2_ok = verify_sequinclusion(&proof2);
        bool quotes_match = hash_matrix_eq(&proof1.quote, &proof2.quote);

        if (inc1_ok && inc2_ok && quotes_match) {
            if (args->proof_count == 1) {
                printf("verified transclusion\n");
            } else {
                printf("verified transclusion %d\n", i + 1);
            }
        } else {
            if (args->proof_count == 1) {
                printf("failed transclusion. reason: ");
            } else {
                printf("failed transclusion %d. reason: ", i + 1);
            }
            if (!inc1_ok) printf("(sequinclusion 1 failed)");
            if (!inc2_ok) printf("(sequinclusion 2 failed)");
            if (!quotes_match) printf("(transclusion failed)");
            printf("\n");
        }
    }
}

void do_verify_transclusion_from_source(Args* args) {
    bool one_indexed = get_indexing_mode(args->index_base);

    int total_sources = args->file_count + args->text_count;
    if (total_sources != 2) {
        fprintf(stderr, "error: transclusion verify requires exactly 2 sources (got %d)\n", total_sources);
        exit(1);
    }

    if (args->index_count != 2) {
        fprintf(stderr, "error: number of sources (-f/-t) must match number of indices (-i). got %d sources, %d indices\n", total_sources, args->index_count);
        exit(1);
    }

    uint8_t* contents1;
    uint8_t* contents2;
    size_t len1, len2;
    bool free1 = false, free2 = false;

    // Get first source
    if (args->file_count > 0) {
        contents1 = read_file(args->files[0], &len1);
        if (!contents1) exit(1);
        free1 = true;
    } else {
        contents1 = (uint8_t*)args->texts[0];
        len1 = strlen(args->texts[0]);
    }

    // Get second source
    if (args->file_count == 2) {
        contents2 = read_file(args->files[1], &len2);
        if (!contents2) {
            if (free1) free(contents1);
            exit(1);
        }
        free2 = true;
    } else if (args->file_count == 1 && args->text_count == 1) {
        contents2 = (uint8_t*)args->texts[0];
        len2 = strlen(args->texts[0]);
    } else {
        contents2 = (uint8_t*)args->texts[1];
        len2 = strlen(args->texts[1]);
    }

    SequinclusionProof proof1, proof2;
    if (compute_sequinclusion_proof(contents1, len1, args->indices[0].start, args->indices[0].end, one_indexed, &proof1) != 0) {
        printf("failed transclusion (error computing proof 1)\n");
        if (free1) free(contents1);
        if (free2) free(contents2);
        exit(1);
    }

    if (compute_sequinclusion_proof(contents2, len2, args->indices[1].start, args->indices[1].end, one_indexed, &proof2) != 0) {
        printf("failed transclusion (error computing proof 2)\n");
        if (free1) free(contents1);
        if (free2) free(contents2);
        exit(1);
    }

    bool inc1_ok = verify_sequinclusion(&proof1);
    bool inc2_ok = verify_sequinclusion(&proof2);
    bool quotes_match = hash_matrix_eq(&proof1.quote, &proof2.quote);

    if (inc1_ok && inc2_ok && quotes_match) {
        printf("verified transclusion\n");
    } else {
        printf("failed transclusion. reason: ");
        if (!inc1_ok) printf("(sequinclusion 1 failed)");
        if (!inc2_ok) printf("(sequinclusion 2 failed)");
        if (!quotes_match) printf("(transclusion failed)");
        printf("\n");
        if (free1) free(contents1);
        if (free2) free(contents2);
        exit(1);
    }

    if (free1) free(contents1);
    if (free2) free(contents2);
}

void do_singlehash(Args* args) {
    // If no files or texts, read from stdin
    if (args->file_count == 0 && args->text_count == 0) {
        uint8_t buffer[1024 * 1024];  // 1MB buffer
        size_t total = 0;
        size_t n;
        while ((n = fread(buffer + total, 1, sizeof(buffer) - total, stdin)) > 0) {
            total += n;
            if (total >= sizeof(buffer)) break;
        }
        hash_matrix_t h = bromberg_hash(buffer, total);
        char hex[65];
        bromberg_to_hex(h, hex);
        printf("%s\n", hex);
        return;
    }

    // Hash each file
    for (int i = 0; i < args->file_count; i++) {
        size_t len;
        uint8_t* contents = read_file(args->files[i], &len);
        if (!contents) exit(1);
        hash_matrix_t h = bromberg_hash(contents, len);
        char hex[65];
        bromberg_to_hex(h, hex);
        printf("%s\n", hex);
        free(contents);
    }

    // Hash each text
    for (int i = 0; i < args->text_count; i++) {
        hash_matrix_t h = bromberg_hash((uint8_t*)args->texts[i], strlen(args->texts[i]));
        char hex[65];
        bromberg_to_hex(h, hex);
        printf("%s\n", hex);
    }
}

int main(int argc, char** argv) {
    bromberg_init();

    Args args = parse_args(argc, argv);

    switch (args.command) {
        case CMD_SEQUINCLUSION_GENERATE:
            do_generate_sequinclusion(&args);
            break;
        case CMD_SEQUINCLUSION_VERIFY:
            do_verify_sequinclusion(&args);
            break;
        case CMD_TRANSCLUSION_GENERATE:
            do_generate_transclusion(&args);
            break;
        case CMD_TRANSCLUSION_VERIFY:
            if (args.file_count > 0 || args.text_count > 0) {
                do_verify_transclusion_from_source(&args);
            } else {
                do_verify_transclusion_proofs(&args);
            }
            break;
        case CMD_SINGLEHASH:
            do_singlehash(&args);
            break;
        default:
            fprintf(stderr, "error: no command specified\n");
            return 1;
    }

    return 0;
}
