import requests
from typing import Callable,Any
from dataclasses import dataclass, field

from rich.console import Console
import datetime
import getpass
@dataclass
class Agent:
    system_prompt: str = "answer in french only"
    model: str = "qwen/qwen3.5-9b"
    base_url: str = "http://127.0.0.1:1234/v1"
    api_key:  str = field(default="NO_API_KEY", repr=False)
    contexts: dict[str, Callable[[],str]] = field(default_factory=dict)
    messages: list[dict[str,Any]] = field(default_factory=list)

    def __post_init__(self) -> None:
        self.base_url = self.base_url.rstrip('/')
    
    def context(self, func : Callable[[],str]) -> Callable[[],str]:
        self.contexts[func.__name__] = func
        return func

    def chat(self, user_message: str) -> str:
        self.messages.append({"role":"user", "content":user_message})

        context_content = "\n\n".join(
            f"\n<{n}>{fn()}<\n"
            for n,fn in self.contexts.items()
        )

        prefix: list[dict[str, Any]] = [
            {"role": "system", "content": self.system_prompt},
            # IS this correct ??? system twice??
            {"role": "system", "content": context_content},
        ]

        url = f"{self.base_url}/chat/completions"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        r = requests.post(
            url,
            headers=headers,
            json={"model": self.model, "messages": prefix + self.messages},
            timeout=300,
        )
        r.raise_for_status()
        data = r.json()
        choices = data.get("choices")

        if not choices:
            raise RuntimeError("Model response missing choices")

        message = choices[0].get("message")
        if message is None:
            raise RuntimeError("response missing message")

        response = message.get("content") or ""
        self.messages.append({"role":"assistant","content": response})
        return response
def main() -> None:
    agent = Agent(
        model="qwen/qwen3.5-9b",
        system_prompt="Tell me a rodney dangerfield joke at end of response",
    )
    @agent.context 
    def user_context() -> str:
        return(
            # f"Current data and time: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
            f"Current data and time: {datetime.datetime.now().strftime('%Y-%m- %H:%M:%S')}\n"
            f"Current User: {getpass.getuser()}\n"
        )

    console = Console()
    with console.status("[dim]Thinking...[/dim]",spinner="arc"):
        response = agent.chat('What time is it and who am i')
    console.print(f"[blue]You:[/blue] what time is it and who am i?")
    console.print(f"[blue]Ai Assistant:[/blue] {response}")
    
 
if __name__ ==  '__main__':
    main()