# -*- coding: utf-8 -*-
"""Sisko AI: FinKing - Fast Finance AI
Using distilgpt2 for Ultra-Fast CPU Inference (Debug Mode)
"""

import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr
import warnings
warnings.filterwarnings('ignore')

print("🚀 SISKO CAPITAL: FINKING - FAST MODE SETUP")

# ===============================
# AUTHENTICATION
# ===============================
print("\n[0] Authenticating with Hugging Face...")
from huggingface_hub import login
hf_token = os.environ.get("HF_TOKEN")
if hf_token:
    login(token=hf_token)
    print("✓ Authenticated with Hugging Face")
else:
    print("⚠ HF_TOKEN not set, continuing...")

# ===============================
# MODEL LOADING (distilgpt2: Ultra-Fast & Reliable)
# ===============================
print("\n[1] Loading Base Model (distilgpt2)...")
model_name = "distilgpt2"

try:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    
    if torch.cuda.is_available():
        model = model.cuda()
        print("✓ Model on GPU")
    else:
        print("✓ Model on CPU (distilgpt2 is optimized for this)")
    
    print("✓ Base Model Loaded")
except Exception as e:
    print(f"⚠ Model loading error: {e}")
    raise

# ===============================
# INFERENCE ENGINE
# ===============================
print("\n[2] Loading Inference Engine...")

def sisko_query(user_query, max_tokens=20):
    """
    Generate response using distilgpt2.
    """
    try:
        prompt = f"Question: {user_query}\nAnswer:"
        
        inputs = tokenizer(prompt, return_tensors="pt")
        device = next(model.parameters()).device
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_tokens,
                temperature=0.7,
                top_p=0.9,
                do_sample=True
            )
        
        full_resp = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return full_resp.strip() if full_resp else "Unable to generate response."
    
    except Exception as e:
        return f"Error: {str(e)[:80]}"

print("✓ Inference Engine Ready (distilgpt2)")

# ===============================
# GRADIO UI
# ===============================
print("\n[3] Launching Gradio UI...")

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
    gr.Markdown(
        """
        # 🤖 Sisko AI: FinKing
        ### AI-Powered Investing for Superior Returns
        **Powered by distilgpt2 (Debug Mode)**
        
        Annual Return: **27%** | Sharpe Ratio: **0.82** | Volatility: **12%**
        """
    )
    
    with gr.Row():
        prompt_input = gr.Textbox(label="Ask me anything", placeholder="What is 2+2?")
        output = gr.Textbox(label="Response")
    
    submit_btn = gr.Button("Submit")
    submit_btn.click(fn=sisko_query, inputs=prompt_input, outputs=output)
    
    gr.Examples(
        examples=["What is 2+2?", "Tell me about AAPL", "Bitcoin outlook?"],
        inputs=prompt_input,
        outputs=output,
        fn=sisko_query,
        cache_examples=False
    )
    
    gr.Markdown(
        """
        ---
        **Contact:** sisko@duck.com | UEN: T25LL0878B | 177 Tanjong Rhu Road, Singapore
        """
    )

demo.launch(share=True)

print("\n✅ SISKO AI LIVE - FAST MODE ACTIVE (distilgpt2)")