#!/usr/bin/env python3

from dataclasses import dataclass
import re
import sys
from typing import Dict, List, Optional, Set, Tuple

HEADER = re.compile(r'===([^=]+)===')
DECL = re.compile(r'([A-Za-z0-9_]+) \(([^)]*)\)')
LINK = re.compile(r'\+.*')

def err(msg):
    print(f'\x1b[31m{msg}\x1b[39m')

@dataclass
class Variable:
    name: str
    current_value: str
    allowed_values: Set[str]

@dataclass
class Choice:
    action: str
    result: Dict[str, str]

    def apply(self, state: Dict[str, Variable]):
        for var, val in self.result.items():
            state[var].current_value = val

@dataclass
class Storylet:
    conditions: Optional[Dict[str, str]]
    text: List[str]
    choices: Optional[List[Choice]]

    def prompt(self) -> Optional[Choice]:
        for ln in ''.join(self.text).strip().split('\n'):
            print('|', ln)

        if not self.choices:
            return None

        for i, chc in enumerate(self.choices):
            print(f'\x1b[1m  {i}: {chc.action}\x1b[0m')

        x = None
        total = len(self.choices)
        if total > 1:
            choice_text = f'0-{total-1}'
        else:
            choice_text = '0'

        while True:
            raw = input(f'Your choice ({choice_text}): ')
            try:
                x = int(raw)
                if x >= 0 and x < len(self.choices):
                    break
            except ValueError:
                pass
            err(f'Invalid choice: {raw}. Choose a number {choice_text}')

        return self.choices[x]

    def for_state(self, state: Dict[str, Variable]) -> Tuple[int, 'Storylet']:
        if self.conditions is None:
            return 0, self

        matching = 0
        for var, desired in self.conditions.items():
            if state[var].current_value != desired:
                return 0, self
            matching += 1

        return matching, self

def parse_state(st: str) -> Dict[str, str]:
    conditions = {}
    st = st.strip().lstrip('[').rstrip(']')
    for thing in st.split(','):
        key, value = thing.split(':')
        conditions[key.strip()] = value.strip()
    return conditions

def parse_file(f):
    state = {}
    chunks = []
    curr_chunk = None

    for ln in f:
        if (match := HEADER.match(ln)):
            stuff = match.group(1).strip()
            if stuff == 'intro':
                conditions = None
            else:
                conditions = parse_state(stuff)
            curr_chunk = Storylet(
                conditions=conditions,
                text=[],
                choices=[],
            )
            chunks.append(curr_chunk)

        elif curr_chunk:
            if (match := LINK.match(ln)):
                if ln.strip() == '+ END':
                    curr_chunk.choice = None
                else:
                    choice, result = ln.strip('+').split('->')
                    result = parse_state(result)
                    choice = Choice(
                        action=choice.strip(),
                        result=result,
                    )
                    curr_chunk.choices.append(choice)
            else:
                curr_chunk.text.append(ln)

        elif (match := DECL.match(ln)):
            name = match.group(1)
            states = match.group(2).split()
            curr = [s.strip('+') for s in states if s.startswith('+')][0]
            allowed = {s.strip('+') for s in states}
            state[name] = Variable(
                name=name,
                current_value=curr,
                allowed_values=allowed,
            )

    return state, chunks

def debug_state(state):
    line = ', '.join(f'{k}: {v.current_value}' for k, v in state.items())
    print(f'\x1b[2m    current state: {line}\x1b[0m')

def run_story(state, chunks):
    current_chunk = [ c for c in chunks if c.conditions is None ]
    if not current_chunk:
        raise("Cannot find initial storylet!")

    current_chunk = current_chunk[0]
    while True:
        choice = current_chunk.prompt()
        if choice is None: return
        choice.apply(state)
        # debug_state(state)
        ranked = sorted([c.for_state(state) for c in chunks], key=lambda x: x[0])
        current_chunk = ranked[-1][1]

def main(path):
    with open(path) as f:
        state, chunks = parse_file(f)
    run_story(state, chunks)

if __name__ == '__main__':
    main(sys.argv[1])