import anthropic
import os
from dotenv import load_dotenv
import yaml
from openai import OpenAI
from huggingface_hub import InferenceClient


load_dotenv(override=True)
os.environ['OPENAI_API_KEY'] = os.getenv('OPENAI_API_KEY')
os.environ['ANTHROPIC_API_KEY'] = os.getenv('ANTHROPIC_API_KEY')
google_api_key = os.getenv('GOOGLE_API_KEY')
deepseek_api_key = os.getenv('DEEPSEEK_API_KEY')
grok_api_key = os.getenv("XAI_API_KEY")

with open('config/var_dev.yaml', 'r') as f:
    config = yaml.safe_load(f)
    

# Logging stuff
if config["local"]:
        hf_token = os.getenv('HF_TOKEN') # Sign in to HuggingFace Hub
else:
        hf_token =  None
#huggingface_hub.login(hf_token)
# Initialize clients
openai = OpenAI()
deepseek_api= OpenAI(
        api_key=deepseek_api_key, 
        base_url="https://api.deepseek.com"
        )
gemini_api = OpenAI(
        api_key=google_api_key, 
        base_url="https://generativelanguage.googleapis.com/v1beta/openai/"
        )
grok_api = OpenAI(api_key=grok_api_key, base_url="https://api.x.ai/v1")
claude = anthropic.Anthropic()
client = InferenceClient()     # For HuggingFace Inference API

LLM_MODEL = None

def user_prompt_for(user_msg):
    return f"""
Trading strategy description:
\"\"\"{user_msg}\"\"\"

Task:
- Convert the description into executable Python code.
- Use only the library {config["api_fin"]}.
- Respond only with valid Python code, following Python best practices
"""

example_1= f''' 
# User prompt:
# "Go long when the 10-period SMA crosses above the 100-period SMA,
#  and exit when the 10-period SMA crosses below the 100-period SMA."


# Generated Python code:

import backtrader as bt
class SmaCross(bt.Strategy):
    """
    Simple moving average crossover strategy.
    Buy when fast SMA crosses above slow SMA.
    Sell when fast SMA crosses below slow SMA.
    """
    params = dict(pfast=10, pslow=100)

    def __init__(self):
        self.sma_fast = bt.ind.SMA(period=self.p.pfast)
        self.sma_slow = bt.ind.SMA(period=self.p.pslow)
        self.crossover = bt.ind.CrossOver(self.sma_fast, self.sma_slow)

    def next(self):
        if not self.position:
            if self.crossover > 0:  # Golden cross
                self.buy()
        elif self.crossover < 0:   # Death cross
            self.close()

# Initialize Cerebro
cerebro = bt.Cerebro()
cerebro.addstrategy(SmaCross, pfast=10, pslow=100)
'''
list_of_pyclasses = [example_1]
system_message = f'''
You are a financial assistant specialized in transforming natural language descriptions of trading strategies into clean, production-ready Python code.

Guidelines:
- Use only the library {config["api_fin"]}.
- Always create a class with the abreviation of the strategy with the form `NameOfStrategy(bt.Strategy)`.
- Implement strategy logic in `__init__` (indicators/signals) and `next()` (trade execution).
- Implement the strategy for this intervall of time {config["interval"]} 
- Finish with initializing the strategy in Cerebro:
  cerebro = bt.Cerebro()
  cerebro.addstrategy(MyStrategy, param1=value, param2=value)
- Keep code minimal, clear, and follow Java best practices (PEP8, clear naming, modularity).
- If a strategy cannot be implemented with {config["api_fin"]}, respond with: "Unable to implement with {config["api_fin"]}."
- If used any addional libraries, add it in the code: import MyUsedLibrary
- If you don't know the answer, just say that you don't know, don't try to make up an answer.


Example(s) of transformation from user prompt of Python code: \n
'''
for pyclass in list_of_pyclasses:
  system_message += pyclass
# Messages in Openai format 
def messages_for(user_msg):
    return [
        {"role": "system", "content": system_message},
        {"role": "user", "content": user_prompt_for(user_msg)}
    ]
    

def stream_llms(user_msg, typ_llm="gpt"): 
    """ Stream responses from different LLMs based on user message and selected model. """  
    global LLM_MODEL
    llm_key = typ_llm.lower()
    model_overrides = {
        "deepseek": ("DEEPSEEK_MODEL", config.get("deepseek_model")),
        "gimini": ("GEMINI_MODEL", config.get("gemini_model")),
        "grok4": ("GROK4_MODEL", config.get("grok4_model")),
        "claude": ("CLAUDE_MODEL", config.get("claude_model")),
        "gpt": ("OPENAI_MODEL", config.get("openai_model")),
    }
    if llm_key not in model_overrides:
        raise ValueError("Unknown model")
    env_var, default_model = model_overrides[llm_key]
    if not default_model:
        raise KeyError(f"Missing default model configuration for '{llm_key}'")
    selected_model = os.getenv(env_var) or default_model
    LLM_MODEL = selected_model
    messages = messages_for(user_msg)

    if llm_key == "deepseek": 
        stream = deepseek_api.chat.completions.create(
            model=selected_model,
            messages=messages,
            stream=True
        )
    elif llm_key == "gimini":
        stream = gemini_api.chat.completions.create(
            model=selected_model,
            messages=messages,
            stream=True
        )
    elif llm_key == "grok4":
        stream = grok_api.chat.completions.create(
            model=selected_model,
            messages=messages,
            stream= True
        )
    elif llm_key == "claude":
        stream =  claude.messages.stream(
        model=selected_model,
        max_tokens=2000,
        system=messages[0]['content'],
        messages=[messages[1]],
    )
    elif llm_key == "gpt": 
        stream = openai.chat.completions.create(model=selected_model, messages=messages, stream=True)
    else:
        raise ValueError("Unknown model")
    
    reply = f"# Model {LLM_MODEL}\n"
    
    if typ_llm.lower() == "claude":
        with stream as stream_clde:
            for fragment in stream_clde.text_stream:
                reply += fragment
                yield reply.replace("```python\n","").replace("```","") 

    else:
        for chunk in stream:
            if chunk and chunk.choices:
                fragment = chunk.choices[0].delta.content or ""
                reply += fragment
                yield reply.replace("```python\n","").replace("```","") 



def stream_manager(user_msg, model):
    """ Streaming manager for different LLMs based on user message and selected model. """
    result = stream_llms(user_msg, model)
    for stream_so_far in result:
        yield stream_so_far