import os, requests, gradio as gr
from smolagents import CodeAgent, GoogleSearchTool, OpenAIServerModel
from smolagents.tools import Tool
from dotenv import load_dotenv

load_dotenv()
ETHERSCAN_API_KEY = os.getenv("ETHERSCAN_API_KEY")
MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
SERPAPI_API_KEY = os.getenv("SERPAPI_API_KEY")

##Smolagent
model = OpenAIServerModel(
    model_id="mistral-medium-latest",
    api_base="https://api.mistral.ai/v1",
    api_key=MISTRAL_API_KEY
)

search_tool = GoogleSearchTool(provider="serpapi")

agent = CodeAgent(
    tools=[search_tool],
    model=model,
    add_base_tools=False
)

##################
###Multi Agent####
##################

###Survey Agent
def classify_risk(*answers):
    score_map = {
        # Q1
        "Short-term (0–2 yrs)": 1, "Medium-term (2–5 yrs)": 2, "Long-term (5+ yrs)": 3,
        # Q2
        "Very uncomfortable": 1, "Somewhat uncomfortable": 2, "Comfortable": 3,
        # Q3
        "Capital preservation": 1, "Moderate growth": 2, "High growth": 3,
        # Q4
        "Beginner": 1, "Intermediate": 2, "Advanced": 3,
        # Q5
        "No": 1, "Maybe": 2, "Yes": 3,
        # Q6
        "<10%": 1, "10–30%": 2, ">30%": 3,
        # Q7
        "Sell immediately": 1, "Hold": 2, "Buy more": 3,
        # Q8
        "Stablecoins": 1, "BTC/ETH": 2, "Meme/Altcoins": 3,
        # Q9
        "Daily": 3, "Weekly": 2, "Rarely": 1,
        # Q10
        "Safe & slow": 1, "Balanced": 2, "Fast & risky": 3
    }

    total = sum(score_map.get(ans, 0) for ans in answers)
    avg = total / 10

    if avg <= 1.6:
        profile = "Conservative"
    elif avg <= 2.3:
        profile = "Moderate"
    else:
        profile = "Aggressive"

    return profile, f"🧠 Your Risk Profile: **{profile}**\n\nScore: {total}/30"

###Portfolio Import Agent
ERC20_TOKENS = {
    "USDC": "0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48",
    "DAI": "0x6b175474e89094c44da98b954eedeac495271d0f",
    "USDT": "0xdac17f958d2ee523a2206206994597c13d831ec7",
    "LINK": "0x514910771af9ca656af840dff83e8264ecf986ca"
}

def import_portfolio(address):
    if not address.startswith("0x") or len(address) != 42:
        return [], "❌ Invalid Ethereum address format.", None

    portfolio = []

    eth_url = f"https://api.etherscan.io/api?module=account&action=balance&address={address}&tag=latest&apikey={ETHERSCAN_API_KEY}"
    try:
        eth_resp = requests.get(eth_url).json()
        eth_balance = int(eth_resp["result"]) / 1e18
        portfolio.append({"asset": "ETH", "balance": eth_balance, "type": "Large-cap"})
    except Exception as e:
        return [], f"Failed to fetch ETH balance: {e}", None

    for token_name, contract in ERC20_TOKENS.items():
        token_url = f"https://api.etherscan.io/api?module=account&action=tokenbalance&contractaddress={contract}&address={address}&tag=latest&apikey={ETHERSCAN_API_KEY}"
        try:
            token_resp = requests.get(token_url).json()
            token_balance = int(token_resp["result"]) / 1e6
            if token_balance > 0:
                token_type = "Stablecoin" if token_name in ["USDC", "DAI"] else "Altcoin"
                portfolio.append({"asset": token_name, "balance": token_balance, "type": token_type})
        except Exception:
            continue

    if not portfolio:
        return [], "No assets found.", None

    # For display and downstream
    display = [[p["asset"], p["balance"]] for p in portfolio]
    return display, "✅ Portfolio fetched.", portfolio

###Analysis Agent
def analyze_portfolio(portfolio, risk_profile):
    if not portfolio:
        return "❌ No portfolio data to analyze."

    total_balance = sum(p["balance"] for p in portfolio)
    if total_balance == 0:
        return "❌ Portfolio is empty."

    # Volatility estimate
    vol_score = 0
    for p in portfolio:
        if p["type"] == "Altcoin":
            vol_score += 3 * (p["balance"] / total_balance)
        elif p["type"] == "Large-cap":
            vol_score += 2 * (p["balance"] / total_balance)
        elif p["type"] == "Stablecoin":
            vol_score += 1 * (p["balance"] / total_balance)

    if vol_score <= 1.5:
        vol_level = "Low"
    elif vol_score <= 2.2:
        vol_level = "Moderate"
    else:
        vol_level = "High"

    # Concentration check
    top_asset_pct = max(p["balance"] / total_balance for p in portfolio)
    concentration = "High" if top_asset_pct > 0.5 else "Moderate" if top_asset_pct > 0.3 else "Low"

    # Diversification
    unique_types = set(p["type"] for p in portfolio)
    diversification = "High" if len(unique_types) >= 3 else "Moderate" if len(unique_types) == 2 else "Low"

    # Risk alignment
    alignment = "Aligned"
    if risk_profile == "Conservative" and vol_level == "High":
        alignment = "Not Aligned"
    elif risk_profile == "Aggressive" and vol_level == "Low":
        alignment = "Under Risked"

    report = f"""
📊 **Portfolio Risk Analysis**
- **Volatility Level:** {vol_level}
- **Top Asset Concentration:** {concentration}
- **Diversification Level:** {diversification}
- **Risk Alignment with Profile ({risk_profile}):** {alignment}
"""

    return report.strip(), report.strip()

##Expert Consulting Agent
def chat_with_advisor(user_message, chat_history, risk_profile, analysis_report):
    if not risk_profile or not analysis_report:
        return chat_history, chat_history + [["⚠️", "Complete survey & analysis first."]]

    prompt = (
        f"Risk Profile:\n{risk_profile}\n\n"
        f"Portfolio Analysis:\n{analysis_report}\n\n"
        f"User: {user_message}\n"
        f"Advisor:"
    )

    try:
        response = agent.run(prompt)
    except Exception as e:
        response = f"❌ Error during agent processing: {e}"

    chat_history.append([user_message, response])
    return chat_history, chat_history
    
#################
### GRADIO UI ###
#################

###Risk Survey
with gr.Blocks() as demo:
    # ✨ New Title Added Here
    gr.Markdown("# Web3WealthManagement")
    gr.Markdown("---") # Optional: Adds a horizontal line for separation

    risk_profile_state = gr.State()  # 🧠 persist profile here
    portfolio_state = gr.State()
    chat_state = gr.State([])  # stores message history
    analysis_state = gr.State()  # stores the analysis string

    gr.Markdown("## 🧪 Step 1: Crypto Risk Survey") # Changed to Step 1 for clarity
    with gr.Row():
        q1 = gr.Radio(["Short-term (0–2 yrs)", "Medium-term (2–5 yrs)", "Long-term (5+ yrs)"], label="1. What's your investment horizon?")
        q2 = gr.Radio(["Very uncomfortable", "Somewhat uncomfortable", "Comfortable"], label="2. How do you feel about a 20% drop?")
        q3 = gr.Radio(["Capital preservation", "Moderate growth", "High growth"], label="3. What's your goal?")
        q4 = gr.Radio(["Beginner", "Intermediate", "Advanced"], label="4. Experience in crypto investing?")
        q5 = gr.Radio(["No", "Maybe", "Yes"], label="5. Would you buy more after a 30% dip?")

    with gr.Row():
        q6 = gr.Radio(["<10%", "10–30%", ">30%"], label="6. How much of your savings would you invest in crypto?")
        q7 = gr.Radio(["Sell immediately", "Hold", "Buy more"], label="7. What would you do in a crash?")
        q8 = gr.Radio(["Stablecoins", "BTC/ETH", "Meme/Altcoins"], label="8. Preferred investment type?")
        q9 = gr.Radio(["Daily", "Weekly", "Rarely"], label="9. How often do you check your portfolio?")
        q10 = gr.Radio(["Safe & slow", "Balanced", "Fast & risky"], label="10. Preferred return strategy?")


    submit_btn = gr.Button("Submit Survey")
    result = gr.Markdown()

    submit_btn.click(
    classify_risk,
    inputs=[q1, q2, q3, q4, q5, q6, q7, q8, q9, q10],
    outputs=[risk_profile_state, result]
)

### Wallet Import
    gr.Markdown("## 📥 Step 2: Import Your Portfolio")
    gr.Markdown("You may try with a sample wallet or enter your own Ethereum address below:")
    wallet_input = gr.Textbox(label="Enter Ethereum Wallet Address")
    with gr.Row():
        vitalik_btn = gr.Button("Use Vitalik's Wallet")
        cuban_btn = gr.Button("Use Mark Cuban's Wallet")
    import_btn = gr.Button("Import Portfolio")
    portfolio_output = gr.Dataframe(headers=["Asset", "Balance"], row_count=5)
    import_status = gr.Markdown()
    import_btn.click(
    import_portfolio,
    inputs=[wallet_input],
    outputs=[portfolio_output, import_status, portfolio_state]
)

### Risk Analysis
    gr.Markdown("## 🔍 Step 3: Analyze Portfolio Risk")
    analyze_btn = gr.Button("Run Risk Analysis")
    analysis_output = gr.Markdown()

    analyze_btn.click(
        analyze_portfolio,
        inputs=[portfolio_state, risk_profile_state],
        outputs=[analysis_output, analysis_state]
    )

### Expert Followup
    gr.Markdown("## 💬 Step 4: Expert Consultation Chat")


    chatbot = gr.Chatbot()
    user_input = gr.Textbox(label="Ask the advisor", placeholder="e.g. How do I reduce my portfolio risk?")
    send_btn = gr.Button("Send")

    send_btn.click(
        chat_with_advisor,
        inputs=[user_input, chat_state, risk_profile_state, analysis_state],
        outputs=[chatbot, chat_state]
    )

    vitalik_btn.click(lambda: "0xd8dA6BF26964aF9D7eEd9e03E53415D37aA96045", outputs=wallet_input)
    cuban_btn.click(lambda: "0xa679c6154b8d4619Af9F83f0bF9a13A680e01eCf", outputs=wallet_input)

demo.launch()