import importlib
import json
import re
import subprocess
import sys

from ollama import chat, ChatResponse


__MODEL__ = 'qwen2.5-coder:7b'
#__MODEL__ = 'llama3.2:latest'
#__MODEL__ = 'qwen3:8b'
#__MODEL__ = 'mistral:7b'

__OPTIONS__ = dict(
    temperature=0.8,
    #'repeat_penalty': 1.2
)

__PREFIX__ = 'code-'


def send_message_to_ollama(messages, ins):
    """Envoie de nouvelles instructions à ollama."""

    ins = ins.lstrip().rstrip()

    print('--------------------')
    print('<<')
    print(ins)
    print('<<')

    msg = {
        'role': 'user',
        'content': ins
    }

    messages.append(msg)

    with open('messages.json', 'w') as fd:
        json.dump(messages, fd)

    resp: ChatResponse = chat(
        model=__MODEL__,
        options=__OPTIONS__,
        messages=messages,
    )

    #messages.append(resp.message)

    msg = {
        'role': resp.message.role,
        'content': resp.message.content
    }

    messages.append(msg)

    print('--------------------')
    print('>>')
    print(resp.message.content)
    print('>>')
    print('--------------------')

    with open('messages.json', 'w') as fd:
        json.dump(messages, fd)

    return resp.message.content


def dump_smali_code(data, index):
    """Extrait le code Smali fournit par ollama."""

    pattern = r'```(.*?)```'
    code_blocks = re.findall(pattern, data, re.DOTALL)

    for blk in code_blocks:

        if blk.startswith('smali'):
            blk = blk[len('smali'):]

        with open(__PREFIX__ + '%04u.smali' % index, 'w') as fd:
            fd.write(blk)

        data = data.replace(blk, '')
        print(data)

        print('Code written!')

        break


def compile_file(filename):
    """Lance une compilation avec smali."""

    cmd = 'java -jar smali-3.0.9-fat.jar a ' + filename + ' -o test.dex'

    process = subprocess.Popen(
        cmd.split(' '),
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
        text=True
    )

    # process.returncode sera toujours 0...
    _, stderr = process.communicate()

    return stderr


def extract_errors(data):
    """Liste toutes les erreurs renvoyées par l'assemblage."""

    found = []

    pat = re.compile("^.*" + __PREFIX__ + "\d+.smali\[(\d+),(\d+)] (.*)$")

    lines = data.split('\n')

    for l in lines:

        match = pat.match(l)

        if match:

            found.append({
                'line': int(match.group(1)),
                'col': int(match.group(2)),
                'msg': match.group(3)
            })

    return found


def check_code(index, config):
    """Génère au besoin une nouvelle requête pour corriger le code fourni."""

    filename = __PREFIX__ + '%04u.smali' % index

    stderr = compile_file(filename)

    errors = extract_errors(stderr)

    print('[i] Errors? %d' % len(errors))

    next_msg = []

    with open(filename, 'r') as fd:
        content = fd.read().split('\n')

    for e in errors:

        assert(e['line'] > 0)

        next_msg.append('- at line %d column %d: %s (erroneous line content : "%s" )' \
                        % (e['line'], e['col'], e['msg'], content[e['line'] - 1].lstrip()))

    if len(next_msg) == 0:
        next_msg = ''

    else:
        next_msg = \
            config.__ERROR_PROLOGUE__.lstrip('\n') \
            + '\n'.join(next_msg) \
            + config.__ERROR_EPILOGUE__

    return next_msg


if __name__ == '__main__':
    """Point d'entrée."""

    argc = len(sys.argv)

    if argc < 2:
        print('Usage: %s <config number> [--check]' % sys.argv[0])
        sys.exit(1)

    config = importlib.import_module('config%02u' % int(sys.argv[1]))

    if argc == 2:

        messages = []

        if config.__SYSTEM__PROMPT__:

            msg = {
                'role': 'system',
                'content': config.__SYSTEM__PROMPT__
            }

            messages.append(msg)

        response = send_message_to_ollama(messages, config.__FIRST_INSTRUCTIONS__)

        counter = 0
        dump_smali_code(response, counter)

        while True:

            instructions = check_code(counter, config)

            if len(instructions) == 0:
                break

            response = send_message_to_ollama(messages, instructions)

            counter += 1
            dump_smali_code(response, counter)

    else:

        filename = sys.argv[2]

        stderr = compile_file(filename)

        errors = extract_errors(stderr)

        print('%s: %d' % (filename, len(errors)))

        sys.exit(0 if len(errors) == 0 else 1)