#!/bin/python3
import os
from pathlib import Path
from openai import OpenAI
from openai.types.responses import EasyInputMessageParam as Message
from typing import Optional, List, Dict, TypeVar, TypeAlias, Callable, Tuple, Literal
import argparse
import re
import subprocess
import readline
import json
from dataclasses import dataclass
from actions import *

T = TypeVar("T")
A = TypeVar("A")
B = TypeVar("B")

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

DiscussionState: TypeAlias = Tuple[Message, ...]
DiscussionM: TypeAlias = Tuple[str, DiscussionState]
QueryMethod: TypeAlias = Callable[
    [str, DiscussionState], DiscussionM
]  # DiscussionM -> DiscussionM if multiple arguments are interpreted as product types


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

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


def mk_direct_query(client: OpenAI, model_name: str) -> QueryMethod:
    def direct_query(
        query: str,
        state: DiscussionState,
        role: Literal["user", "system", "developer"] = "user",
    ) -> DiscussionM:
        logger.info(f"send: {query}")

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

        res = client.responses.create(model=model_name, input=list(state))

        logger.info(f"receive: {res.output_text}")

        return (
            res.output_text,
            state + ({"role": "assistant", "content": res.output_text},),
        )

    return direct_query


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) = config.query_method(
        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)",
    )
    parser.add_argument(
        "--endpoint",
        type=str,
        default=os.environ.get(
            "LLM_API_ENDPOINT", "https://llm-gateway.k8s.aalto.fi/api/v1"
        ),
        help="The API endpoint for the LLM",
    )
    parser.add_argument(
        "--key",
        type=str,
        default=os.environ.get("LLM_API_KEY"),
        help="The key for the LLM API endpoint",
    )
    parser.add_argument(
        "--model",
        type=str,
        default="RedHatAI/gemma-4-31B-it-FP8-Dynamic",
        help="The model name to use",
    )
    parser.add_argument(
        "--trust-level",
        type=int,
        default=2,
        help="The trust level of the agent (default: 2)",
    )
    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=args.trust_level,
        query_method=mk_direct_query(
            OpenAI(base_url=args.endpoint, api_key=args.key), args.model
        ),
    )

    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.
""",
        },
    )

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

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

        print(response)


if __name__ == "__main__":
    main()