from __future__ import annotations
import os
from typing import Optional

# Providers are optional; we import lazily


class LLMClient:
	def __init__(self) -> None:
		self.provider = os.getenv("LLM_PROVIDER", "openai").lower()
		self.openai_key = os.getenv("OPENAI_API_KEY")
		self.anthropic_key = os.getenv("ANTHROPIC_API_KEY")
		self.gemini_key = os.getenv("GEMINI_API_KEY")
		self._openai_client = None
		self._anthropic_client = None
		self._gemini_model = None

		# Optional per-agent Gemini keys (fallback to default if missing)
		self._agent_keys = {
			"cv": os.getenv("GEMINI_API_KEY_CV") or self.gemini_key,
			"cover": os.getenv("GEMINI_API_KEY_COVER") or self.gemini_key,
			"chat": os.getenv("GEMINI_API_KEY_CHAT") or self.gemini_key,
			"parser": os.getenv("GEMINI_API_KEY_PARSER") or self.gemini_key,
			"match": os.getenv("GEMINI_API_KEY_MATCH") or self.gemini_key,
			"tailor": os.getenv("GEMINI_API_KEY_TAILOR") or self.gemini_key,
		}

		# Preload if configured
		if self.provider == "openai" and self.openai_key:
			try:
				from openai import OpenAI
				self._openai_client = OpenAI(api_key=self.openai_key)
			except Exception:
				self._openai_client = None
		elif self.provider == "anthropic" and self.anthropic_key:
			try:
				import anthropic
				self._anthropic_client = anthropic.Anthropic(api_key=self.anthropic_key)
			except Exception:
				self._anthropic_client = None
		elif self.provider == "gemini" and self.gemini_key:
			# We will lazily configure per-call to support per-agent keys
			try:
				import google.generativeai as genai  # noqa: F401
			except Exception:
				self._gemini_model = None

	@property
	def enabled(self) -> bool:
		if self.provider == "openai":
			return self._openai_client is not None
		if self.provider == "anthropic":
			return self._anthropic_client is not None
		if self.provider == "gemini":
			# If we have at least one usable key, consider enabled
			return any([self.gemini_key] + list(self._agent_keys.values()))
		return False

	def generate(self, system_prompt: str, user_prompt: str, model: Optional[str] = None, max_tokens: int = 1200, agent: Optional[str] = None) -> str:
		# Fallback behavior if no provider configured
		if not self.enabled:
			text = (system_prompt + "\n\n" + user_prompt)[: max_tokens * 4]
			return text

		provider = self.provider
		if provider == "openai":
			return self._generate_openai(system_prompt, user_prompt, model or os.getenv("LLM_MODEL", "gpt-4o-mini"), max_tokens)
		if provider == "anthropic":
			return self._generate_anthropic(system_prompt, user_prompt, model or os.getenv("LLM_MODEL", "claude-3-5-sonnet-latest"), max_tokens)
		if provider == "gemini":
			return self._generate_gemini(system_prompt, user_prompt, model or os.getenv("LLM_MODEL", "gemini-1.5-flash"), max_tokens, agent=agent)
		# Unknown provider fallback
		return (system_prompt + "\n\n" + user_prompt)[: max_tokens * 4]

	def _generate_openai(self, system_prompt: str, user_prompt: str, model: str, max_tokens: int) -> str:
		try:
			response = self._openai_client.chat.completions.create(
				model=model,
				messages=[
					{"role": "system", "content": system_prompt},
					{"role": "user", "content": user_prompt},
				],
				temperature=0.4,
				max_tokens=max_tokens,
			)
			return response.choices[0].message.content.strip()
		except Exception:
			return (system_prompt + "\n\n" + user_prompt)[: max_tokens * 4]

	def _generate_anthropic(self, system_prompt: str, user_prompt: str, model: str, max_tokens: int) -> str:
		try:
			msg = self._anthropic_client.messages.create(
				model=model,
				max_tokens=max_tokens,
				system=system_prompt,
				messages=[{"role": "user", "content": user_prompt}],
				temperature=0.4,
			)
			# Anthropic returns a list of content blocks
			parts = []
			for b in msg.content:
				if hasattr(b, "text"):
					parts.append(b.text)
				elif isinstance(b, dict) and b.get("type") == "text":
					parts.append(b.get("text", ""))
			return "\n".join(p for p in parts if p).strip() or (system_prompt + "\n\n" + user_prompt)[: max_tokens * 4]
		except Exception:
			return (system_prompt + "\n\n" + user_prompt)[: max_tokens * 4]

	def _generate_gemini(self, system_prompt: str, user_prompt: str, model: str, max_tokens: int, agent: Optional[str] = None) -> str:
		try:
			import google.generativeai as genai
			# Resolve API key per agent if provided
			api_key = self.gemini_key
			if agent:
				# Normalize agent to known keys
				norm = agent.lower()
				if norm == "general":
					norm = "chat"
				api_key = self._agent_keys.get(norm, self.gemini_key)
			# Configure and call
			genai.configure(api_key=api_key)
			model_instance = genai.GenerativeModel(model)
			prompt = system_prompt + "\n\n" + user_prompt
			resp = model_instance.generate_content(prompt)
			text = getattr(resp, "text", None)
			if not text and hasattr(resp, "candidates") and resp.candidates:
				text = resp.candidates[0].content.parts[0].text
			return (text or prompt)[: max_tokens * 4]
		except Exception:
			return (system_prompt + "\n\n" + user_prompt)[: max_tokens * 4]


llm = LLMClient()