#!/bin/python3
import os
from pathlib import Path
from openai import OpenAI
from typing import Optional, List, Dict, TypeVar, TypeAlias, Callable
import argparse
import re
import subprocess
import readline
import json
from dataclasses import dataclass
from actions import *

T = TypeVar("T")

import logging

logger = logging.getLogger(__name__)
logging.basicConfig(filename="thoughts.log", level=logging.INFO)

JSON: TypeAlias = Dict[str, "JSON"] | List["JSON"] | str | int | float | bool | None


@dataclass
class AgentConfig:
    """Configuration for agents."""

    workdir: Path
    actions: List["Action"]
    trustlevels: Dict[str, int]
    trustlevel: int
    search_padding: int = 200


client = OpenAI(
    api_key=os.environ["AALTO_LLM_KEY"],
    base_url="https://llm-gateway.k8s.aalto.fi/api/v1",
)


def direct_query(query: str):
    logger.info(f"send: {query}")
    res = client.responses.create(
        model="RedHatAI/gemma-4-31B-it-FP8-Dynamic", input=query
    )
    logger.info(f"receive: {res.output_text}")
    return res.output_text


def chat(content, state=(), query_method=direct_query, role="user"):

    state = state + ({"role": role, "content": content},)

    response = query_method(state)

    state = state + ({"role": "assistant", "content": response},)

    return (response, state)


def first(l: List[T], p: Callable[[T], bool]) -> Optional[T]:
    return next((x for x in l if p(x)), None)


def query_agent(query: str, config: AgentConfig, state=()):

    state = state + ({"role": "user", "content": query},)

    (res, state) = chat(
        f"Based on the current state, what should you do next?\n"
        "Respond JSON indicating an action, or respond in natural language if you are finished.",
        state,
    )

    # Replace all fenced JSON blocks

    replacements = []
    result_jsons: List[JSON] = []
    for m in re.finditer(r"^```json\s*(.*?)\s*^```", res, re.DOTALL | re.MULTILINE):
        json_str = m.group(1)
        try:
            action = json.loads(json_str)
            action_name = action["action"]
        except (json.JSONDecodeError, KeyError):
            continue

        f = first(config.actions, lambda f: f.name == action_name)
        if f is not None:
            desc = f.formatter(action)
            if (
                config.trustlevels[action_name] > config.trustlevel
                and not input(f"Perform {desc}? [y/N] ") == "y"
            ):
                result_json = {"status": "error", "reason": "Permission denied by user"}
            else:
                result_json = f(action, config)
            result_jsons.append(result_json)
            replacements.append((m.start(), m.end(), desc))

    replaced_res = ""
    i = 0
    for start, end, action in replacements:
        replaced_res += res[i:start] + f"[performed '{action}']\n"
        i = end
    replaced_res += res[i:]

    if len(replacements) > 0:
        (res, state) = query_agent(json.dumps(result_jsons), config, state)
        return (replaced_res + "\n" + res, state)
    else:
        return (replaced_res, state)


def main():
    parser = argparse.ArgumentParser(description="Autonomous Coding Agent")
    parser.add_argument(
        "--workdir",
        type=Path,
        default=Path.cwd(),
        help="The directory the agent has access to (default: current directory)",
    )
    args = parser.parse_args()

    config = AgentConfig(
        workdir=args.workdir,
        actions=[ls, read, write, tldr, shell, mypy, webgrep],
        trustlevels={
            "ls": 0,
            "read": 0,
            "tldr": 0,
            "mypy": 0,
            "webgrep": 1,
            "write": 2,
            "shell": 3,
        },
        trustlevel=2,
    )

    state = (
        {
            "role": "system",
            "content": f"""You are an autonomous coding agent that can help with programming tasks.
You can explore repositories, read files, search for information, and make changes to code.

You have access to the following actions, which you should embed in JSON markdown blocks "```json ACTION ```" in your responses.
{"\n".join(map(lambda f: "* " + f.description, config.actions))}

After reading the relevant files, remember to analyze the task at hand and break it into steps.

Once you are finished with your task, give a response with none of the above JSON commands.
""",
        },
    )

    print(state)

    while True:
        query = input("> ").strip()

        (response, state) = query_agent(query, config, state)

        print(response)


if __name__ == "__main__":
    main()