import os
import uuid
import requests
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

TS_HOST        = os.getenv("TS_HOST")
TS_USERNAME    = os.getenv("TS_USERNAME")
TS_SECRET_KEY  = os.getenv("TS_SECRET_KEY")
TOKEN_VALIDITY = int(os.getenv("TOKEN_VALIDITY", "300"))   # seconds
MCP_URL        = os.getenv("MCP_URL", "https://agent.thoughtspot.app/bearer/mcp")



#  Generate TS_AUTH_TOKEN using secret_key (same as curl)

def generate_ts_token():
    

    if not TS_HOST or not TS_USERNAME or not TS_SECRET_KEY:
        raise ValueError("Missing TS_HOST, TS_USERNAME, or TS_SECRET_KEY in .env")

    url = f"https://{TS_HOST}/api/rest/2.0/auth/token/full"

    payload = {
        "username": TS_USERNAME,
        "validity_time_in_sec": TOKEN_VALIDITY,
        "auto_create": False,
        "secret_key": TS_SECRET_KEY,
    }

    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json"
    }

    print("\n Generating ThoughtSpot token...")
    resp = requests.post(url, json=payload, headers=headers)

    print(" Status:", resp.status_code)
    print(" Raw response:", resp.text[:300])

    if resp.status_code != 200:
        raise Exception(f"Token generation failed: {resp.text}")

    token = resp.json().get("token")
    if not token:
        raise Exception("No 'token' field returned from ThoughtSpot")

    print(" TS_AUTH_TOKEN generated successfully.")
    return token




# 2. JSON-RPC Helper for MCP
def _mcp_call(method: str, params: dict | None, token: str, session_id: str | None, rpc_id: int):
    
    headers = {
        "Authorization": f"Bearer {token}",
        "x-ts-host": TS_HOST,
        "Content-Type": "application/json",
        "Accept": "application/json, text/event-stream",
    }

    if session_id:
        headers["Mcp-Session-Id"] = session_id

    payload = {
        "jsonrpc": "2.0",
        "id": str(rpc_id),
        "method": method,
        "params": params or {},
    }

    return requests.post(MCP_URL, headers=headers, json=payload, timeout=45)



# 3.  MCP interface: initialize → tools/list → tools/call
def ask_thoughtspot_via_mcp(user_question: str) -> dict:
    

    #  Always generate a fresh TS_AUTH_TOKEN
    token = generate_ts_token()

    print("\n Calling ThoughtSpot MCP")
    print("MCP_URL:", MCP_URL)
    print("TS_HOST:", TS_HOST)

    # initialize (NO session ID)
    init_payload = {
        "protocolVersion": "2024-11-05",
        "capabilities": {},
        "clientInfo": {"name": "rowan-hybrid-bot", "version": "1.0.0"},
    }

    r1 = _mcp_call("initialize", init_payload, token, None, rpc_id=1)
    print("\n initialize():", r1.status_code)
    print(r1.text[:250])

    if r1.status_code != 200:
        return {"ok": False, "error": f"initialize failed ({r1.status_code}): {r1.text[:300]}"}

    #  list available tools (WITH session ID)
    session_id = str(uuid.uuid4())
    r2 = _mcp_call("tools/list", {}, token, session_id, rpc_id=2)
    print("\n tools/list():", r2.status_code)
    print(r2.text[:250])

    if r2.status_code != 200:
        return {"ok": False, "error": f"tools/list failed ({r2.status_code}): {r2.text[:300]}"}

    # Extract tools
    try:
        tools = [t.get("name") for t in r2.json().get("result", {}).get("tools", [])]
    except Exception:
        tools = []

    if not tools:
        return {"ok": False, "error": "No MCP tools available for this user/token."}

    # Prefer getAnswer
    tool_name = "getAnswer" if "getAnswer" in tools else tools[0]

    #  Step 3 — Call the tool
    args = {"question": user_question}
    r3 = _mcp_call("tools/call", {"name": tool_name, "arguments": args}, token, session_id, rpc_id=3)

    print("\n tools/call():", r3.status_code)
    print(r3.text[:250])

    if r3.status_code != 200:
        return {"ok": False, "tool": tool_name, "error": r3.text[:300]}

    # Successful call
    try:
        raw = r3.json()
        return {"ok": True, "tool": tool_name, "raw": raw}
    except Exception:
        return {"ok": True, "tool": tool_name, "answer": r3.text}