import os
import json
import time
import uuid
from typing import Dict, Any, Optional
import requests

class LLMClient:
    def __init__(self, cfg: Dict[str, Any], logs_dir: str):
        self.cfg = cfg
        self.logs_dir = logs_dir
        os.makedirs(self.logs_dir, exist_ok=True)

    def _log(self, role: str, content: str, uid: str):
        path = os.path.join(self.logs_dir, f"{time.strftime('%Y%m%d-%H%M%S')}-{uid}-{role}.md")
        with open(path, "w", encoding="utf-8") as f:
            f.write(content)

    def chat(self, system: str, user: str) -> str:
        uid = str(uuid.uuid4())[:8]
        backend = self.cfg.get("backend", "ollama")
        if backend == "ollama":
            host = self.cfg["ollama"]["host"]
            model = self.cfg["ollama"]["model"]
            payload = {
                "model": model,
                "messages": [
                    {"role": "system", "content": system},
                    {"role": "user", "content": user},
                ],
                "stream": False,
            }
            self._log("system", system, uid)
            self._log("user", user, uid)
            r = requests.post(f"{host}/v1/chat/completions", json=payload, timeout=600)
            r.raise_for_status()
            data = r.json()
            content = data["choices"][0]["message"]["content"]
            self._log("assistant", content, uid)
            return content

        elif backend == "openai":
            import openai  # requires openai>=1.0
            openai.api_key = os.environ.get("OPENAI_API_KEY")
            if not openai.api_key:
                raise RuntimeError("Missing OPENAI_API_KEY")
            model = self.cfg["openai"]["model"]
            messages = [
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ]
            self._log("system", system, uid)
            self._log("user", user, uid)
            client = openai.OpenAI()
            resp = client.chat.completions.create(model=model, messages=messages)
            content = resp.choices[0].message.content
            self._log("assistant", content, uid)
            return content

        else:
            raise ValueError(f"Unknown backend: {backend}")