#ifndef PROBOL_RUNTIME_H
#define PROBOL_RUNTIME_H

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <termios.h>
#include <sys/ioctl.h>

#define MAX_DIST_STATES 32

// ---------------------------------------------------------------------
// Core probabilistic structures
// ---------------------------------------------------------------------
typedef struct {
    double value;
    double probability;
} ProbPair;

typedef struct {
    ProbPair states[MAX_DIST_STATES];
    int state_count;
} ProbTensor;

static inline void init_prob_tensor(ProbTensor* t) {
    t->state_count = 0;
    memset(t->states, 0, sizeof(t->states));
}

static inline void add_state(ProbTensor* t, double val, double prob) {
    if (t->state_count < MAX_DIST_STATES) {
        t->states[t->state_count].value = val;
        t->states[t->state_count].probability = prob;
        t->state_count++;
    }
}

static inline void cct_add(ProbTensor* dest, ProbTensor* a, ProbTensor* b) {
    ProbTensor temp;
    init_prob_tensor(&temp);
    for (int i = 0; i < a->state_count; i++) {
        for (int j = 0; j < b->state_count; j++) {
            double combined_val = a->states[i].value + b->states[j].value;
            double combined_prob = a->states[i].probability * b->states[j].probability;
            int found = 0;
            for (int k = 0; k < temp.state_count; k++) {
                if (temp.states[k].value == combined_val) {
                    temp.states[k].probability += combined_prob;
                    found = 1;
                    break;
                }
            }
            if (!found) {
                add_state(&temp, combined_val, combined_prob);
            }
        }
    }
    *dest = temp;
}

static inline double cct_measure(ProbTensor* t) {
    if (t->state_count == 0) return 0.0;
    int best_idx = 0;
    double max_p = -1.0;
    for (int i = 0; i < t->state_count; i++) {
        if (t->states[i].probability > max_p) {
            max_p = t->states[i].probability;
            best_idx = i;
        }
    }
    return t->states[best_idx].value;
}

// ---------------------------------------------------------------------
// Terminal I/O
// ---------------------------------------------------------------------
static inline void prob_term_puts(const char* s) {
    if (s) write(STDOUT_FILENO, s, strlen(s));
}

static inline void prob_term_putc(char c) {
    write(STDOUT_FILENO, &c, 1);
}

static inline char prob_term_getc(void) {
    char c = 0;
    (void)read(STDIN_FILENO, &c, 1);
    return c;
}

static inline void prob_term_clear(void) {
    prob_term_puts("\033[2J\033[H");
}

static inline void prob_term_color(int code) {
    char seq[12];
    snprintf(seq, sizeof(seq), "\033[3%dm", code & 7);
    prob_term_puts(seq);
}

static inline void prob_term_goto(int x, int y) {
    char seq[32];
    snprintf(seq, sizeof(seq), "\033[%d;%dH", y, x);
    prob_term_puts(seq);
}

static inline void prob_term_size(int* w, int* h) {
    struct winsize ws;
    if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &ws) == 0) {
        *w = ws.ws_col;
        *h = ws.ws_row;
    } else {
        *w = 80; *h = 24;
    }
}

static struct termios _probol_orig_termios;
static inline void prob_term_raw_enable(void) {
    tcgetattr(STDIN_FILENO, &_probol_orig_termios);
    struct termios raw = _probol_orig_termios;
    raw.c_iflag &= ~(BRKINT | ICRNL | INPCK | ISTRIP | IXON);
    raw.c_oflag &= ~(OPOST);
    raw.c_cflag |= (CS8);
    raw.c_lflag &= ~(ECHO | ICANON | IEXTEN | ISIG);
    raw.c_cc[VMIN] = 0;
    raw.c_cc[VTIME] = 1;
    tcsetattr(STDIN_FILENO, TCSAFLUSH, &raw);
}

static inline void prob_term_raw_disable(void) {
    tcsetattr(STDIN_FILENO, TCSAFLUSH, &_probol_orig_termios);
}

static inline void print_prob_tensor(const ProbTensor* t) {
    printf("ProbTensor (%d states):\n", t->state_count);
    for (int i = 0; i < t->state_count; i++) {
        printf("  [%.4f] -> %.6f\n", t->states[i].value, t->states[i].probability);
    }
}

#endif