import os
import json
import logging
from typing import List, Optional
from fastapi import FastAPI, HTTPException, Header
from pydantic import BaseModel
from sentence_transformers import CrossEncoder
from huggingface_hub import InferenceClient

# --- LOGGING ---
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")
logger = logging.getLogger("Brain")

app = FastAPI()

# --- CONFIG ---
EXPECTED_TOKEN = os.getenv("SERVICE_SECRET", "default_secret")
HF_TOKEN = os.getenv("HF_TOKEN")

# --- MODELS ---
logger.info("⏳ Loading CrossEncoder...")
matcher_model = CrossEncoder("cross-encoder/stsb-distilroberta-base")
logger.info("✅ CrossEncoder Ready")

llm_client = InferenceClient(model="Qwen/Qwen2.5-72B-Instruct", token=HF_TOKEN)

# --- INPUT SCHEMAS ---
class MatchRequest(BaseModel):
    target: str
    candidates: List[str]

class MarketData(BaseModel):
    rsi: float
    macd: float
    atr_pct: float
    trend: str
    price: float
    bias: str
    prob_up: float

class AnalyzeRequest(BaseModel):
    question: str
    context: str = "Crypto 15m Market"
    data: Optional[MarketData] = None # ¡Ahora recibimos datos!

# --- ENDPOINTS ---

@app.get("/")
def health():
    return {"status": "Brain V2 is Active 🧠", "matcher": "Loaded"}

@app.post("/match")
def match_market(req: MatchRequest, x_service_token: str = Header(None)):
    if x_service_token != EXPECTED_TOKEN: raise HTTPException(401, "Unauthorized")
    
    pairs = [[req.target, c] for c in req.candidates]
    scores = matcher_model.predict(pairs)
    best_idx = scores.argmax()
    best_score = float(scores[best_idx])
    
    if best_score > 0.4: # Bajamos umbral para captar más
        return {"match": req.candidates[best_idx], "score": best_score}
    return {"match": None, "score": best_score}

@app.post("/analyze")
def analyze_market(req: AnalyzeRequest, x_service_token: str = Header(None)):
    if x_service_token != EXPECTED_TOKEN: raise HTTPException(401, "Unauthorized")
    
    # Construimos un Prompt Rico con los datos técnicos
    tech_summary = "No technical data provided."
    if req.data:
        tech_summary = (
            f"STRATEGY SIGNAL: {req.data.bias} (Conf: {req.data.prob_up:.0%})\n"
            f"INDICATORS: RSI={req.data.rsi:.1f} | MACD={req.data.macd:.4f} | Trend={req.data.trend}\n"
            f"VOLATILITY (ATR): {req.data.atr_pct:.3f}% | Price: {req.data.price}"
        )

    logger.info(f"🤔 Analyzing: {req.question} | Bias: {req.data.bias if req.data else 'N/A'}")

    prompt = f"""
    You are a Senior Crypto Risk Officer for a high-frequency trading bot.
    
    MARKET: "{req.question}"
    
    TECHNICAL ANALYSIS:
    {tech_summary}
    
    TASK:
    1. Check if the Market Title implies a "Trap" (e.g., weird rules, Elon Musk tweets, specific candle closes).
    2. Validate if the Strategy Signal ({req.data.bias if req.data else 'N/A'}) makes sense given the indicators.
       - Example: Buying BULLISH when RSI is > 80 is RISKY (Overbought).
       - Example: Selling BEARISH when RSI is < 20 is RISKY (Oversold).
       - Example: Trading when ATR is > 0.5% is RISKY (Too volatile).
    
    OUTPUT JSON ONLY:
    {{"verdict": "SAFE" or "TRAP", "reason": "concise explanation"}}
    """
    
    try:
        response = llm_client.chat_completion(
            messages=[{"role": "user", "content": prompt}],
            max_tokens=150, temperature=0.1
        )
        content = response.choices[0].message.content.replace("```json", "").replace("```", "").strip()
        return json.loads(content)
    except Exception as e:
        logger.error(f"LLM Error: {e}")
        return {"verdict": "SAFE", "reason": "LLM Failed, defaulting to Safe"}