import os
from typing import Optional
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv

load_dotenv()

def get_llm(
    temperature: float = 0.7,
    model_name_override: Optional[str] = None,
    provider: Optional[str] = None,
):
    # provider: "openai" | None (None = use USE_OPENAI env)
    """
    Factory for LLM. Optional provider overrides env:
    - provider="openai" -> OpenAI (OPENAI_API_KEY, OPENAI_MODEL_NAME)
    - provider="openrouter" -> OpenRouter (OPENROUTER_API_KEY, OPENROUTER_MODEL_NAME)
    - provider=None -> use USE_OPENAI env (true -> OpenAI, else OpenRouter).
    Other provider strings raise ValueError.
    """
    if provider is not None:
        provider_norm = str(provider).strip().lower()
        if provider_norm not in {"openai", "openrouter"}:
            raise ValueError("provider must be one of: 'openai', 'openrouter'")
        use_openai = provider_norm == "openai"
    else:
        use_openai = (os.getenv("USE_OPENAI") or "false").strip().lower() == "true"

    if use_openai:
        api_key = os.getenv("OPENAI_API_KEY")
        model_name = model_name_override or os.getenv("OPENAI_MODEL_NAME", "gpt-4o-mini")
        if not api_key:
            raise ValueError("OPENAI_API_KEY not found in environment.")
        
        return ChatOpenAI(
            openai_api_key=api_key,
            model_name=model_name,
            temperature=temperature
        )
    else:
        api_key = os.getenv("OPENROUTER_API_KEY")
        model_name = model_name_override or os.getenv("OPENROUTER_MODEL_NAME", "tngtech/deepseek-r1t2-chimera:free")
        if not api_key:
            raise ValueError("OPENROUTER_API_KEY not found in environment.")
        
        return ChatOpenAI(
            openai_api_key=api_key,
            openai_api_base="https://openrouter.ai/api/v1",
            model_name=model_name,
            temperature=temperature,
            extra_body={
                "HTTP-Referer": "https://github.com/vaibhavk/profiling",
                "X-Title": "User Profiling Engine"
            }
        )