#include <ctype.h>
#include <stddef.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <stdio.h>

#include "lex.h"
#include "utils.h"



typedef struct {
    char* str;
    size_t pos;
} RawLine;


String collect_identifier(RawLine* line) {
    String id_buf = make_String(256);
    while (line->str[line->pos] == '_' || isalnum(line->str[line->pos])) {
        push_to_String(&id_buf, line->str[line->pos++]);
    }
    return id_buf;
}



TokenVec lex_level(RawLine* line, char until, RegionVec* reg_arena);


void expect_char(RawLine* line, char expected) {
    if (line->str[line->pos] != expected) {
        exit(45);     // TODO: nicer handling
    }
    line->pos++; 
}


MaybeChunk lex_expansion_chunk(RawLine* line, RegionVec* reg_arena, QuotationMode qmode) {
    line->pos++;    // Consume $
    if (line->str[line->pos] == '(') {
        line->pos++;
        if (line->str[line->pos] == '(') {
            exit(1);    
            // TODO: this (and potentially the other parsers must be custom, we need to recognize arithm-internal parens from the terminating `))` )
        

            TokenVec inner_toks = lex_level(line, ')', reg_arena);
            expect_char(line, ')');
            expect_char(line, ')');
            Region* reg_ptr = push_to_RegionVec(
                reg_arena, 
                (Region){.mode = ARITHM, .tokens = inner_toks}
            );
            return (MaybeChunk){
                .some = true,
                .value = (Chunk){
                    .qmode = qmode, 
                    .kind = REGION_EXP, 
                    .as.reg_ptr = reg_ptr
                },
            };
        } else {
            TokenVec inner_toks = lex_level(line, ')', reg_arena);
            expect_char(line, ')');
            Region* reg_ptr = push_to_RegionVec(
                reg_arena, 
                (Region){.mode = CMDS, .tokens = inner_toks}
            );
            return (MaybeChunk) {
                .some = true,
                .value = (Chunk){
                    .qmode = qmode,
                    .kind = REGION_EXP,
                    .as.reg_ptr = reg_ptr,
                }
            };
        }
    } else {
        String id_name = collect_identifier(line);
        if (id_name.count > 0) {
            return (MaybeChunk){
                .some = true,
                .value = (Chunk){
                    .qmode = qmode,
                    .kind = VAR_EXP,
                    .as.varname = inner(&id_name),
                }
            };
        } 
    }
    return (MaybeChunk){.some = false};
} 


void lex_double_quoted_section(RawLine* line, Word* coll, RegionVec* reg_arena) {
    line->pos++;
    String accum = make_String(256);
    while (line->str[line->pos] != '\"') {
        if (line->str[line->pos] == '\0') exit(43);

        MaybeChunk exp_chunk = {.some = false};

        if (line->str[line->pos] == '$') {
            exp_chunk = lex_expansion_chunk(line, reg_arena, DOUBLE_Q);
            if (!exp_chunk.some) {
                    push_to_String(&accum, '$');
            }            
        } else {
            push_to_String(&accum, line->str[line->pos++]);
        }
        // TODO: other special cases beyond $
        
        if (exp_chunk.some) {
            if (accum.count > 0) {
                push_to_Word(coll, (Chunk){
                    .qmode = DOUBLE_Q, 
                    .kind = LITERAL, 
                    .as.literal_str = inner(&accum)
                });
                accum = make_String(256);
            }
            push_to_Word(coll, exp_chunk.value);
            exp_chunk.some = false;
        }
    }
    line->pos++;
    if (accum.count > 0) {
        push_to_Word(coll, (Chunk){
            .qmode = DOUBLE_Q, 
            .kind = LITERAL, 
            .as.literal_str = inner(&accum)
        });
    }
}


Chunk next_word_unquoted(RawLine* line, RegionVec* reg_arena) {
    String buf = make_String(256);
    for (char c; (c=line->str[line->pos], c != '$' && c != '\0' && c != ' ');  line->pos++) {
        push_to_String(&buf, c);
    }
    return (Chunk){.qmode = UN_Q, .kind = LITERAL, .as.literal_str = inner(&buf)};
}


Chunk collect_single_quoted_section(RawLine* line) {
    line->pos++;    // Consume quote mark
    String chunk_str = make_String(256);
    while (line->str[line->pos] != '\'') {
        if (line->str[line->pos] == '\0') exit(44);
        push_to_String(&chunk_str, line->str[line->pos++]);
    }
    line->pos++;      // Consume closing quote
    return (Chunk){
        .qmode = SINGLE_Q,
        .kind = LITERAL,
        .as.literal_str = inner(&chunk_str),
    };
}


typedef struct {
    TokenVec tokens;
    Word word_germ;
    String unq_liter_germ;
} LexState;


void close_literal_germ(LexState* state) {
   if (state->unq_liter_germ.count > 0) {
        push_to_Word(
            &state->word_germ, 
            (Chunk){
                .qmode = UN_Q, 
                .kind = LITERAL, 
                .as.literal_str = inner(&state->unq_liter_germ)
            }
        );
        state->unq_liter_germ = make_String(256);
    } 
}

void close_word_and_lit_germs(LexState* state) {
    close_literal_germ(state);
    if (state->word_germ.count > 0) {
        push_to_TokenVec(&state->tokens, (Token){.kind = TOK_WORD, .as.word = state->word_germ});
        state->word_germ = make_Word(256);
    }
}

TokenVec lex_level(RawLine* line, char until, RegionVec* reg_arena) {

    LexState state = {
        .tokens = make_TokenVec(256), 
        .word_germ = make_Word(256), 
        .unq_liter_germ = make_String(256)
    };

    while (line->str[line->pos] != until) {


        if (line->str[line->pos] == ' ') {
            close_word_and_lit_germs(&state);
            while (line->str[line->pos] == ' ') line->pos++;
            continue;
        }

        switch (line->str[line->pos]) {
            case '\'':
                push_to_Word(&state.word_germ, collect_single_quoted_section(line));
                break;
            case '\"':
                lex_double_quoted_section(line, &state.word_germ, reg_arena);
                break;
            case '|':
                line->pos++;
                close_word_and_lit_germs(&state);
                push_to_TokenVec(&state.tokens, (Token){.kind = TOK_SPECIAL, .as.spec = PIPE});
                break;
            case '$':
                line->pos++;
                MaybeChunk mchunk = lex_expansion_chunk(line, reg_arena, UN_Q);
                if (mchunk.some) {
                    push_to_Word(&state.word_germ, mchunk.value);
                } else {
                    push_to_String(&state.unq_liter_germ, '$');
                }
                break;
            case '>':
                close_word_and_lit_germs(&state);
                line->pos++;

                switch (line->str[line->pos]) {
                    case '(':
                        line->pos++;
                        TokenVec sub_regtoks = lex_level(line, ')', reg_arena);
                        Region subreg = {.mode = CMDS, .tokens = sub_regtoks};
                        Region* reg_ptr = push_to_RegionVec(reg_arena, subreg);
                        line->pos++;

                        push_to_TokenVec(&state.tokens, (Token){.kind = TOK_PROCSUB_OUT, .as.procsub_region = reg_ptr});
                        break;
                    case '>':
                        line->pos++;
                        push_to_TokenVec(&state.tokens, (Token){.kind = TOK_SPECIAL, .as.spec = WRITE_REDIR_APPEND});
                        break;
                    default:
                        push_to_TokenVec(&state.tokens, (Token){.kind = TOK_SPECIAL, .as.spec = WRITE_REDIR_TRUNC});
                }
                break;
            case '<':
                close_word_and_lit_germs(&state);
                line->pos++;
                if (line->str[line->pos] == '(') {
                    line->pos++;
                    TokenVec sub_regtoks = lex_level(line, ')', reg_arena);
                    Region subreg = {.mode = CMDS, .tokens = sub_regtoks};
                    Region* reg_ptr = push_to_RegionVec(reg_arena, subreg);
                    line->pos++;

                    push_to_TokenVec(&state.tokens, (Token){.kind = TOK_PROCSUB_IN, .as.procsub_region = reg_ptr});
                } else {
                    push_to_TokenVec(&state.tokens, (Token){.kind = TOK_SPECIAL, .as.spec = READ_REDIR});
                }
                break;
            default:
                push_to_String(&state.unq_liter_germ, line->str[line->pos++]);
        }
    }
    // TODO: handle until char _and_ EOL
    close_word_and_lit_germs(&state);
    return state.tokens;
}


LexResult lex(char* input) {
    RawLine line = {.str = input, .pos = 0};
    RegionVec reg_arena = make_RegionVec(256);
    TokenVec toks = lex_level(&line, '\0', &reg_arena);
    push_to_TokenVec(&toks, (Token){.kind = TOK_SPECIAL, .as.spec = EOL});
    return (LexResult){.top_tokens = toks, .reg_arena = reg_arena};
}