from __future__ import annotations

import asyncio
import os
from uuid import uuid4

from jazzmine.core import (
    AgentBuilder,
    Flow,
    JsonStorage,
    OpenAILLMConfig,
    ServerConfig,
    ToolParameter,
    ToolResponse,
    tool,
)

SANDBOX_NAME = "market"

YAHOO_ALLOWLIST_HOSTS = [
    "query1.finance.yahoo.com",
    "query2.finance.yahoo.com",
    "finance.yahoo.com",
    "download.finance.yahoo.com",
    "fc.yahoo.com",
    "guce.yahoo.com",
    "consent.yahoo.com",
]

TOOL_DEPS = [
    "yfinance==0.2.66",
    "pandas==2.2.3",
    "numpy==2.1.3",
]

@tool(
  description="Fetch a quick market snapshot for one ticker.",
  parameters=[
    ToolParameter("ticker", "str", "Ticker symbol like AAPL.", required=True),
    ToolParameter("period", "str", "5d, 1mo, 3mo, 6mo, 1y", required=False, default="1mo"),
  ],
  dependencies=TOOL_DEPS,
  secrets=[],
  sandbox=SANDBOX_NAME,
)
async def fetch_market_snapshot(ticker: str, period: str = "1mo") -> ToolResponse:
  import pandas as pd
  import yfinance as yf

  symbol = ticker.strip().upper()
  if not symbol:
    return ToolResponse(success=False, message="ticker cannot be empty")

  history = yf.Ticker(symbol).history(period=period, interval="1d", auto_adjust=False, actions=False)
  if history is None or history.empty:
    return ToolResponse(success=False, message=f"No data returned for {symbol}")

  history = history.dropna(subset=["Close"])
  if history.empty:
    return ToolResponse(success=False, message=f"No close price data for {symbol}")

  latest = history.iloc[-1]
  prev_close = float(history["Close"].iloc[-2]) if len(history) >= 2 else None
  close = float(latest["Close"])

  change_abs = None
  change_pct = None
  if prev_close not in (None, 0.0):
    change_abs = close - prev_close
    change_pct = (change_abs / prev_close) * 100.0

  payload = {
    "ticker": symbol,
    "period": period,
    "close": close,
    "open": float(latest["Open"]) if not pd.isna(latest["Open"]) else None,
    "high": float(latest["High"]) if not pd.isna(latest["High"]) else None,
    "low": float(latest["Low"]) if not pd.isna(latest["Low"]) else None,
    "volume": float(latest["Volume"]) if not pd.isna(latest["Volume"]) else None,
    "change_abs": change_abs,
    "change_pct": change_pct,
  }
  return ToolResponse(success=True, data=payload)


@tool(
  description="Compare multiple tickers by total return over a period.",
  parameters=[
    ToolParameter("tickers", "str", "Comma-separated tickers, for example AAPL,MSFT,NVDA.", required=True),
    ToolParameter("period", "str", "1mo, 3mo, 6mo, 1y", required=False, default="6mo"),
  ],
  dependencies=TOOL_DEPS,
  secrets=[],
  sandbox=SANDBOX_NAME,
)
async def compare_ticker_performance(tickers: str, period: str = "6mo") -> ToolResponse:
  import re
  import yfinance as yf

  symbols = [s.strip().upper() for s in re.split(r"[,;\s]+", tickers) if s.strip()]
  symbols = list(dict.fromkeys(symbols))

  if len(symbols) < 2:
    return ToolResponse(success=False, message="Provide at least two tickers")

  data = yf.download(
    tickers=symbols,
    period=period,
    interval="1d",
    auto_adjust=True,
    progress=False,
    threads=False,
  )
  if data is None or data.empty:
    return ToolResponse(success=False, message="No data returned")

  close_df = data["Close"] if hasattr(data.columns, "levels") and "Close" in data.columns.get_level_values(0) else None
  if close_df is None:
    return ToolResponse(success=False, message="Close prices not available")

  close_df = close_df.dropna(how="all")
  if close_df.empty:
    return ToolResponse(success=False, message="No valid close matrix")

  metrics = {}
  for symbol in close_df.columns:
    prices = close_df[symbol].dropna()
    if len(prices) < 2:
      continue
    total_return = float((prices.iloc[-1] / prices.iloc[0]) - 1.0)
    metrics[str(symbol)] = {
      "first_close": float(prices.iloc[0]),
      "last_close": float(prices.iloc[-1]),
      "total_return": total_return,
      "observations": int(len(prices)),
    }

  if len(metrics) < 2:
    return ToolResponse(success=False, message="Could not compute at least two ticker metrics")

  ranked = sorted(metrics.items(), key=lambda item: item[1]["total_return"], reverse=True)
  payload = {
    "period": period,
    "metrics": metrics,
    "rank_by_return": [{"ticker": t, "total_return": m["total_return"]} for t, m in ranked],
    "best_performer": {"ticker": ranked[0][0], "total_return": ranked[0][1]["total_return"]},
    "worst_performer": {"ticker": ranked[-1][0], "total_return": ranked[-1][1]["total_return"]},
  }
  return ToolResponse(success=True, data=payload)

def build_flows() -> list[Flow]:
  return [
    Flow(
      name="market_snapshot",
      description="Fetch latest market snapshot for one ticker.",
      condition="User asks for latest price, move, OHLC, quote, or quick market view.",
      desired_effects=[
        "Call snapshot tool for the requested ticker.",
        "Report close and period clearly.",
      ],
      examples=[
        "Give me a quick snapshot for AAPL",
        "What is the latest price of TSLA?",
        "Show 1 month market snapshot for MSFT",
        "How did NVDA move recently?",
        "Get quote context for AMZN",
      ],
      tools=["fetch_market_snapshot"],
      priority=(1, 0),
    ),
    Flow(
      name="multi_ticker_comparison",
      description="Compare return performance across multiple tickers.",
      condition="User asks to compare two or more tickers by performance.",
      desired_effects=[
        "Call comparison tool with requested symbols.",
        "Summarize best and worst performer clearly.",
      ],
      examples=[
        "Compare AAPL MSFT NVDA over 6 months",
        "Which did better, GOOGL or AMZN this year?",
        "Rank META NFLX DIS by return",
        "Compare TSLA and F",
        "Show me top performer among AAPL, MSFT, NVDA",
      ],
      tools=["compare_ticker_performance"],
      priority=(1, 1),
    ),
    Flow(
      name="methodology_explainer",
      description="Explain market metrics and interpretation guidance.",
      condition="User asks what returns or risk metrics mean.",
      desired_effects=[
        "Explain definitions in plain language.",
        "State that output is informational, not financial advice.",
      ],
      examples=[
        "What does total return mean?",
        "How should I interpret these numbers?",
        "Is this financial advice?",
        "Explain drawdown",
        "What does risk mean here?",
      ],
      tools=[],
      priority=(1, 2),
    ),
  ]

def build_base_agent_builder() -> AgentBuilder:
  api_key = os.environ["MISTRAL_API_KEY"]
  base_url = os.environ.get("MISTRAL_BASE_URL", "https://api.mistral.ai")
  model = os.environ.get("MISTRAL_MODEL", "mistral-medium-latest")

  allowed_ports = {allowed_host: [443] for allowed_host in YAHOO_ALLOWLIST_HOSTS}

  return (
    AgentBuilder(
      name="MarketLite",
      agent_id="market-lite-v1",
      personality=(
        "You are a concise market assistant. Use tools for live numbers, "
        "do not invent market data, and keep outputs informational only."
      ),
    )
    .llm(
      OpenAILLMConfig(
        model=model,
        api_key=api_key,
        base_url=base_url,
        temperature=0.2,
        max_tokens=4000,
      )
    )
    .storage(JsonStorage(path="./market_lite_store.json"))
    .sandbox(
      name=SANDBOX_NAME,
      python_version="3.11",
      timeout_sec=60,
      memory_mb=768,
      cpu_quota=0.75,
      allowed_hosts=YAHOO_ALLOWLIST_HOSTS,
      allowed_ports=allowed_ports,
      default_port=443,
      pip_packages=TOOL_DEPS,
      execution_mode="plan",
      pool_size=1,
    )
    .embeddings('BAAI/bge-small-en-v1.5')
    .memory(
    qdrant_url="http://localhost:6334",
    collection_prefix="market_analyst_pro_v1",
    qdrant_api_key=None,
    qdrant_auto_start=True,
    vector_size=384,
    )
    .flows(build_flows())
    .version("1.0.0")
  )

async def run_interactive() -> None:
  agent, teardown = await build_base_agent_builder().build()
  conversation_id = f"market-lite-{uuid4().hex[:8]}"

  print("MarketLite interactive mode ready")
  print("Commands: /quit")

  try:
    while True:
      user_input = input("\nYou: ").strip()
      if not user_input:
        continue
      if user_input.lower() in {"/quit", "quit", "exit", "q"}:
        break

      result = await agent.chat(
        user_id="market-user",
        conversation_id=conversation_id,
        content=user_input,
      )

      print("\nMarketLite:")
      print(result.response)
      if result.invoked_flows:
        print("[flows]", ", ".join(result.invoked_flows))
      if result.invoked_tools:
        print("[tools]", ", ".join(result.invoked_tools))
  finally:
    await teardown()

async def run_http_server() -> None:
  host = os.environ.get("JAZZMINE_SERVER_HOST", "127.0.0.1")
  port = int(os.environ.get("JAZZMINE_SERVER_PORT", "8010"))
  server_cfg = ServerConfig(host=host, port=port)

  _agent, teardown = await build_base_agent_builder().server(server_cfg).build()

  print("MarketLite server is running")
  print(f"Chat endpoint:          http://{host}:{port}{server_cfg.chat_endpoint}")
  print(f"Stream endpoint:        http://{host}:{port}{server_cfg.chat_endpoint}/stream")
  print(f"Conversations endpoint: http://{host}:{port}{server_cfg.conversations_endpoint}")
  print(f"Health endpoint:        http://{host}:{port}{server_cfg.health_endpoint}")
  print(f"Info endpoint:          http://{host}:{port}{server_cfg.info_endpoint}")

  stop_event = asyncio.Event()
  try:
    await stop_event.wait()
  except (KeyboardInterrupt, asyncio.CancelledError):
    pass
  finally:
    await teardown()

async def main() -> None:
  mode = os.environ.get("MARKET_RUN_MODE", "loop").strip().lower()
  if mode in {"server", "http"}:
    await run_http_server()
  else:
    await run_interactive()


if __name__ == "__main__":
  asyncio.run(main())