#!/bin/python3
import os
from pathlib import Path
import httpx2
from openai import OpenAI, DefaultHttpx2Client
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
    socks5_proxy: Optional[str] = None
    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 int(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)",
    )
    parser.add_argument(
        "--socks5-proxy",
        type=str,
        default=os.environ.get("LLM_SOCKS5_PROXY"),
        help=(
            "Route API queries (and webgrep fetches) through a SOCKS5 proxy, e.g. "
            "socks5://host:1080. Use socks5h:// to resolve DNS through the proxy; "
            "user:pass@ credentials in the URL are supported. "
            "(default: LLM_SOCKS5_PROXY environment variable)"
        ),
    )
    args = parser.parse_args()

    http_client: Optional[httpx2.Client] = (
        DefaultHttpx2Client(proxy=args.socks5_proxy) if args.socks5_proxy else None
    )

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

    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, _) = mk_direct_query(client, args.model, squasher=None)(
                "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,
        ),
        socks5_proxy=args.socks5_proxy,
    )

    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:
        try:
            query = input("> ").strip()
        except KeyboardInterrupt:
            break

        # Don't include this in the try above st. if the user
        # interrupts the agent, it's possible to find out what the
        # program was up to at the moment.
        (response, state) = query_agent(query, config, state)

        print(response)


if __name__ == "__main__":
    main()