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

#include "lex.h"
#include "typeutils.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, ExpansionRegionVec* reg_arena);


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


MaybeChunk lex_expansion_chunk(RawLine* line, ExpansionRegionVec* reg_arena, QuotationMode qmode) {
    line->pos++;    // Consume $
    if (line->str[line->pos] == '(') {
        if (line->str[line->pos] == '(') {
            TokenVec inner_toks = lex_level(line, ')', reg_arena);
            expect_char(line, ')');
            expect_char(line, ')');
            push_to_ExpansionRegionVec(
                reg_arena, 
                (ExpansionRegion){.mode = ARITHM, .tokens = inner_toks}
            );
            ExpansionRegion* reg_ptr = &reg_arena->items[reg_arena->count];  // TODO: maybe have push() return a pointer [in the macro]
            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, ')');
            push_to_ExpansionRegionVec(
                reg_arena, 
                (ExpansionRegion){.mode = CMD_SUB, .tokens = inner_toks}
            );
            ExpansionRegion* reg_ptr = &reg_arena->items[reg_arena->count]; 
            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, ChunkVec* coll, ExpansionRegionVec* reg_arena) {
    line->pos++;
    String accum = make_String(256);
    while (line->str[line->pos] != '\"') {
        if (line->str[line->pos] == '\0') exit(42);

        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_ChunkVec(coll, (Chunk){
                    .qmode = DOUBLE_Q, 
                    .kind = LITERAL, 
                    .as.literal_str = inner(accum)
                });
                accum = make_String(256);
            }
            push_to_ChunkVec(coll, exp_chunk.value);
            exp_chunk.some = false;
        }
    }
    line->pos++;
    if (accum.count > 0) {
        push_to_ChunkVec(coll, (Chunk){
            .qmode = DOUBLE_Q, 
            .kind = LITERAL, 
            .as.literal_str = inner(accum)
        });
    }
}


Chunk next_word_unquoted(RawLine* line, ExpansionRegionVec* 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(42);
        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),
    };
}



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

    TokenVec tokens = make_TokenVec(256);    
    ChunkVec curr_chunks = make_ChunkVec(256);
    String word_buf = make_String(256);

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

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

        switch (line->str[line->pos]) {
            case '\'':
                push_to_ChunkVec(&curr_chunks, collect_single_quoted_section(line));
                break;
            case '\"':
                lex_double_quoted_section(line, &curr_chunks, reg_arena);
                break;
            case '$':
                line->pos++;
                MaybeChunk mchunk = lex_expansion_chunk(line, reg_arena, UN_Q);
                if (mchunk.some) {
                    push_to_ChunkVec(&curr_chunks, mchunk.value);
                } else {
                    push_to_String(&word_buf, '$');
                }
                break;
            default:
                push_to_String(&word_buf, line->str[line->pos++]);
        }
    }

    line->pos++;    // consume until char
    if (word_buf.count > 0) {
        push_to_ChunkVec(
            &curr_chunks, 
            (Chunk){.qmode = UN_Q, .kind = LITERAL, .as.literal_str = inner(word_buf)}
        );
    }
    if (curr_chunks.count > 0) {
        push_to_TokenVec(&tokens, (Token){.kind = TOK_WORD, .as.word = curr_chunks});
    }
    return tokens;
}


typedef struct {
    TokenVec top_tokens;
    ExpansionRegionVec reg_arena;
} LexResult;


TokenVec lex(char* input) {
    RawLine line = {.str = input, .pos = 0};
    ExpansionRegionVec reg_arena = make_ExpansionRegionVec(256);
    return lex_level(&line, '\0', &reg_arena);
}