#!/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 estimate_tokens(state: DiscussionState) -> int:
    """Simple heuristic to estimate tokens in the conversation state."""
    text = "".join(str(m["content"]) for m in state)
    return len(text) // 3.5


def mk_direct_query(
    client: OpenAI,
    model_name: str,
    max_token_count: int = 27500,
    squash_threshold: float = 0.75,
    squasher: Optional[Callable[[DiscussionState], DiscussionState]] = None,
) -> 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},)

        if (
            squasher is not None
            and estimate_tokens(state) > max_token_count * squash_threshold
        ):
            logger.info("Squashing history")
            state = squasher(state)

        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, TypeError):
            continue

        f = first(
            config.actions,
            lambda f: (
                f.name == action_name and all(k in action for k in f.required_keys)
            ),
        )
        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=os.environ.get("LLM_MODEL") or "zai-org/GLM-5.3-Flash",
        help="The model to use",
    )
    parser.add_argument(
        "--min-history",
        type=int,
        default=3,
        help="Minimum amount of most recent messages to keep intact when summarising the history to keep the context short enough",
    )
    parser.add_argument(
        "--model-context-size",
        type=int,
        default=os.environ.get("LLM_CONTEXT_SIZE") or 1048576,
        help="Context size supported by the model.",
    )
    parser.add_argument(
        "--trust-level",
        type=int,
        default=2,
        help="The trust level of the agent (default: 2)",
    )
    args = parser.parse_args()

    client = OpenAI(base_url=args.endpoint, api_key=args.key)

    # We need a version of direct_query that doesn't trigger squashing to use inside the squasher
    # to avoid infinite recursion.
    raw_direct_query = mk_direct_query(client, args.model, squasher=None)

    def squash_history(state: DiscussionState) -> DiscussionState:
        if len(state) <= 1 + args.min_history:
            return (
                state[0],
                {"role": "system", "content": "(history removed)"},
                state[-1],
            )
        else:
            (summary, _) = raw_direct_query(
                "Please summarize the preceding conversation history concisely, preserving all key technical decisions, progress made, and current goals.",
                state[1 : -args.min_history],
                role="system",
            )
            summary_msg: Message = {
                "role": "system",
                "content": f"Previous history summary: {summary}",
            }
            return (state[0], summary_msg) + state[-args.min_history :]

    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(
            client,
            args.model,
            squasher=squash_history,
            max_token_count=args.model_context_size,
        ),
    )

    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()