# app.py — MCP server using DeepSeek via Hugging Face transformers (or fallback)
# - Put this file next to config.py (see example below)
# - It loads the model in LOCAL_MODEL (e.g., a DeepSeek HF checkpoint) via transformers.pipeline
# - If the model cannot be loaded (no transformers / OOM / missing weights), it falls back to a small CPU model or rule-based responder

from mcp.server.fastmcp import FastMCP
from typing import Optional, List, Tuple, Any, Dict
import requests
import os
import gradio as gr
import json
import time
import traceback
import inspect
import re
import logging

# Setup simple logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("mcp_deepseek")

# Optional transformers imports — will attempt; we handle missing gracefully
TRANSFORMERS_AVAILABLE = False
try:
    from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
    TRANSFORMERS_AVAILABLE = True
except Exception as e:
    logger.warning("transformers not available: %s", e)
    TRANSFORMERS_AVAILABLE = False

# ----------------------------
# Load config
# ----------------------------
try:
    from config import (
        CLIENT_ID,
        CLIENT_SECRET,
        REFRESH_TOKEN,
        API_BASE,
        LOCAL_MODEL,      # e.g. "deepseek-ai/deepseek-r1-7b" or smaller/distilled variant
    )
except Exception as e:
    raise SystemExit(
        "Make sure config.py exists with CLIENT_ID, CLIENT_SECRET, REFRESH_TOKEN, API_BASE, LOCAL_MODEL (or set LOCAL_MODEL=None)."
    )

# ----------------------------
# FastMCP init
# ----------------------------
mcp = FastMCP("ZohoCRMAgent")

# ----------------------------
# Analytics / KPI logging (simple local JSON file)
# ----------------------------
ANALYTICS_PATH = "mcp_analytics.json"

def _init_analytics():
    if not os.path.exists(ANALYTICS_PATH):
        base = {"tool_calls": {}, "llm_calls": 0, "last_llm_confidence": None, "created_at": time.time()}
        with open(ANALYTICS_PATH, "w") as f:
            json.dump(base, f, indent=2)

def _log_tool_call(tool_name: str, success: bool = True):
    try:
        with open(ANALYTICS_PATH, "r") as f:
            data = json.load(f)
    except Exception:
        data = {"tool_calls": {}, "llm_calls": 0, "last_llm_confidence": None}
    data["tool_calls"].setdefault(tool_name, {"count": 0, "success": 0, "fail": 0})
    data["tool_calls"][tool_name]["count"] += 1
    if success:
        data["tool_calls"][tool_name]["success"] += 1
    else:
        data["tool_calls"][tool_name]["fail"] += 1
    with open(ANALYTICS_PATH, "w") as f:
        json.dump(data, f, indent=2)

def _log_llm_call(confidence: Optional[float] = None):
    try:
        with open(ANALYTICS_PATH, "r") as f:
            data = json.load(f)
    except Exception:
        data = {"tool_calls": {}, "llm_calls": 0, "last_llm_confidence": None}
    data["llm_calls"] = data.get("llm_calls", 0) + 1
    if confidence is not None:
        data["last_llm_confidence"] = confidence
    with open(ANALYTICS_PATH, "w") as f:
        json.dump(data, f, indent=2)

_init_analytics()

# ----------------------------
# DeepSeek / HF model loader
# ----------------------------
LLM_PIPELINE = None
TOKENIZER = None
LOADED_MODEL_NAME = None

def init_deepseek_model():
    """
    Try to load LOCAL_MODEL via transformers.pipeline.
    Expected LOCAL_MODEL examples:
      - "deepseek-ai/deepseek-r1-7b"  (may require GPU; big)
      - "deepseek-ai/deepseek-r1-3b"  (smaller)
      - "deepseek-ai/deepseek-r1-1.3b" (more likely to load on moderate machines)
    If loading fails, try a fallback small model (distilgpt2 or flan-t5-small if seq2seq).
    """
    global LLM_PIPELINE, TOKENIZER, LOADED_MODEL_NAME

    if not LOCAL_MODEL:
        logger.info("LOCAL_MODEL is None — no local LLM will be used.")
        LLM_PIPELINE = None
        return

    if not TRANSFORMERS_AVAILABLE:
        logger.warning("transformers not installed; cannot load DeepSeek. Falling back to rule-based.")
        LLM_PIPELINE = None
        return

    try:
        tokenizer_name = LOCAL_TOKENIZER or LOCAL_MODEL
        model_name = LOCAL_MODEL
        LOADED_MODEL_NAME = model_name

        # If model looks like seq2seq (T5/flan) use text2text; else causal
        seq2seq_keywords = ["flan", "t5", "seq2seq"]
        if any(k in model_name.lower() for k in seq2seq_keywords):
            TOKENIZER = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True)
            model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
            LLM_PIPELINE = pipeline("text2text-generation", model=model, tokenizer=TOKENIZER)
            logger.info("Loaded seq2seq model: %s", model_name)
        else:
            TOKENIZER = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True)
            model = AutoModelForCausalLM.from_pretrained(model_name)
            LLM_PIPELINE = pipeline("text-generation", model=model, tokenizer=TOKENIZER)
            logger.info("Loaded causal model: %s", model_name)

    except Exception as e:
        logger.error("Failed to load requested model '%s': %s", LOCAL_MODEL, e)
        traceback.print_exc()
        # Try a small CPU-friendly fallback
        fallback = None
        try:
            # prefer an instruction-friendly small model if possible
            fallback = "google/flan-t5-small"
            if "flan" in fallback:
                TOKENIZER = AutoTokenizer.from_pretrained(fallback, use_fast=True)
                model = AutoModelForSeq2SeqLM.from_pretrained(fallback)
                LLM_PIPELINE = pipeline("text2text-generation", model=model, tokenizer=TOKENIZER)
            else:
                TOKENIZER = AutoTokenizer.from_pretrained("distilgpt2", use_fast=True)
                model = AutoModelForCausalLM.from_pretrained("distilgpt2")
                LLM_PIPELINE = pipeline("text-generation", model=model, tokenizer=TOKENIZER)
            LOADED_MODEL_NAME = fallback
            logger.info("Loaded fallback model: %s", fallback)
        except Exception as e2:
            logger.error("Fallback model also failed: %s", e2)
            traceback.print_exc()
            LLM_PIPELINE = None
            LOADED_MODEL_NAME = None

# Initialize model at startup (may take time)
init_deepseek_model()

# ----------------------------
# Rule-based fallback responder
# ----------------------------
def rule_based_response(message: str) -> str:
    msg = (message or "").strip().lower()
    if msg.startswith("create record") or msg.startswith("create contact"):
        return "To create a record, use: create_record MODULE_NAME {\"Field\":\"value\"}"
    if msg.startswith("create_invoice"):
        return "To create invoice: create_invoice {\"customer_id\":\"...\",\"line_items\":[...]} (JSON)"
    if msg.startswith("help") or "what can you do" in msg:
        return "I can run create_record/update_record/delete_record or process local files by pasting their /mnt/data/ path."
    return "(fallback) No local LLM loaded. Use explicit commands like create_record or paste /mnt/data/ path."

# ----------------------------
# LLM wrapper that returns text + confidence (best-effort)
# ----------------------------
def deepseek_generate(prompt: str, max_tokens: int = 256) -> Dict[str, Any]:
    """
    Generate using the loaded pipeline. Returns {'text': str, 'confidence': Optional[float], 'raw': resp}
    """
    if LLM_PIPELINE is None:
        return {"text": rule_based_response(prompt), "confidence": None, "raw": None}

    try:
        out = LLM_PIPELINE(prompt, max_new_tokens=max_tokens)
        text = ""
        # pipeline returns list: [{'generated_text':...}] or [{'generated_text' or 'text'}]
        if isinstance(out, list) and len(out) > 0:
            first = out[0]
            if isinstance(first, dict):
                text = first.get("generated_text") or first.get("generated_text", "") or first.get("text") or str(first)
            else:
                text = str(first)
        else:
            text = str(out)
        _log_llm_call(None)
        return {"text": text, "confidence": None, "raw": out}
    except Exception as e:
        logger.error("LLM generation error: %s", e)
        traceback.print_exc()
        return {"text": rule_based_response(prompt), "confidence": None, "raw": str(e)}

# ----------------------------
# Zoho token refresh & MCP tools (unchanged)
# ----------------------------
def _get_valid_token_headers() -> dict:
    token_url = "https://accounts.zoho.in/oauth/v2/token"
    params = {"refresh_token": REFRESH_TOKEN, "client_id": CLIENT_ID, "client_secret": CLIENT_SECRET, "grant_type": "refresh_token"}
    r = requests.post(token_url, params=params, timeout=20)
    if r.status_code == 200:
        t = r.json().get("access_token")
        return {"Authorization": f"Zoho-oauthtoken {t}"}
    else:
        raise RuntimeError(f"Failed to refresh Zoho token: {r.status_code} {r.text}")

@mcp.tool()
def authenticate_zoho() -> str:
    try:
        _ = _get_valid_token_headers()
        _log_tool_call("authenticate_zoho", True)
        return "Zoho token refreshed (ok)."
    except Exception as e:
        _log_tool_call("authenticate_zoho", False)
        return f"Failed to authenticate: {e}"

@mcp.tool()
def create_record(module_name: str, record_data: dict) -> str:
    try:
        headers = _get_valid_token_headers()
        url = f"{API_BASE}/{module_name}"
        payload = {"data": [record_data]}
        r = requests.post(url, headers=headers, json=payload, timeout=20)
        if r.status_code in (200, 201):
            _log_tool_call("create_record", True)
            return json.dumps(r.json(), ensure_ascii=False)
        _log_tool_call("create_record", False)
        return f"Error creating record: {r.status_code} {r.text}"
    except Exception as e:
        _log_tool_call("create_record", False)
        return f"Exception: {e}"

@mcp.tool()
def get_records(module_name: str, page: int = 1, per_page: int = 200) -> list:
    try:
        headers = _get_valid_token_headers()
        url = f"{API_BASE}/{module_name}"
        r = requests.get(url, headers=headers, params={"page": page, "per_page": per_page}, timeout=20)
        if r.status_code == 200:
            _log_tool_call("get_records", True)
            return r.json().get("data", [])
        _log_tool_call("get_records", False)
        return [f"Error retrieving {module_name}: {r.status_code} {r.text}"]
    except Exception as e:
        _log_tool_call("get_records", False)
        return [f"Exception: {e}"]

@mcp.tool()
def update_record(module_name: str, record_id: str, data: dict) -> str:
    try:
        headers = _get_valid_token_headers()
        url = f"{API_BASE}/{module_name}/{record_id}"
        payload = {"data": [data]}
        r = requests.put(url, headers=headers, json=payload, timeout=20)
        if r.status_code == 200:
            _log_tool_call("update_record", True)
            return json.dumps(r.json(), ensure_ascii=False)
        _log_tool_call("update_record", False)
        return f"Error updating: {r.status_code} {r.text}"
    except Exception as e:
        _log_tool_call("update_record", False)
        return f"Exception: {e}"

@mcp.tool()
def delete_record(module_name: str, record_id: str) -> str:
    try:
        headers = _get_valid_token_headers()
        url = f"{API_BASE}/{module_name}/{record_id}"
        r = requests.delete(url, headers=headers, timeout=20)
        if r.status_code == 200:
            _log_tool_call("delete_record", True)
            return json.dumps(r.json(), ensure_ascii=False)
        _log_tool_call("delete_record", False)
        return f"Error deleting: {r.status_code} {r.text}"
    except Exception as e:
        _log_tool_call("delete_record", False)
        return f"Exception: {e}"

@mcp.tool()
def create_invoice(data: dict) -> str:
    try:
        headers = _get_valid_token_headers()
        url = f"{API_BASE}/invoices"
        r = requests.post(url, headers=headers, json={"data": [data]}, timeout=20)
        if r.status_code in (200, 201):
            _log_tool_call("create_invoice", True)
            return json.dumps(r.json(), ensure_ascii=False)
        _log_tool_call("create_invoice", False)
        return f"Error creating invoice: {r.status_code} {r.text}"
    except Exception as e:
        _log_tool_call("create_invoice", False)
        return f"Exception: {e}"

@mcp.tool()
def process_document(file_path: str, target_module: Optional[str] = "Contacts") -> dict:
    """
    Accepts local path (e.g. /mnt/data/script_zoho_mcp) or URL.
    Per developer instruction we treat the path as the file URL (file://...).
    Replace placeholder OCR logic with your pipeline.
    """
    try:
        if os.path.exists(file_path):
            file_url = f"file://{file_path}"
            extracted = {
                "Name": "ACME Corp (simulated)",
                "Email": "contact@acme.example",
                "Phone": "+91-99999-00000",
                "Total": "1234.00",
                "Confidence": 0.88
            }
            _log_tool_call("process_document", True)
            return {"status": "success", "file": os.path.basename(file_path), "file_url": file_url, "target_module": target_module, "extracted_data": extracted}
        else:
            _log_tool_call("process_document", False)
            return {"status": "error", "error": "file not found", "file_path": file_path}
    except Exception as e:
        _log_tool_call("process_document", False)
        return {"status": "error", "error": str(e)}

# ----------------------------
# Simple command parser (explicit commands in chat)
# ----------------------------
def try_parse_and_invoke_command(text: str):
    text = text.strip()
    m = re.match(r"^create_record\s+(\w+)\s+(.+)$", text, re.I)
    if m:
        module = m.group(1); body = m.group(2)
        try: record_data = json.loads(body)
        except Exception: return "Invalid JSON for record_data"
        return create_record(module, record_data)
    m = re.match(r"^create_invoice\s+(.+)$", text, re.I)
    if m:
        body = m.group(1)
        try: invoice_data = json.loads(body)
        except Exception: return "Invalid JSON for invoice_data"
        return create_invoice(invoice_data)
    m = re.match(r"^(\/mnt\/data\/\S+)$", text)
    if m:
        path = m.group(1); return process_document(path)
    return None

# ----------------------------
# Chat handler that uses DeepSeek generation (or fallback)
# ----------------------------
def deepseek_response(message: str, history: Optional[List[Tuple[str,str]]] = None) -> str:
    history = history or []
    system_prompt = "You are Zoho Assistant. Prefer concise answers. If you want to call a tool, return a JSON object: {\"tool\": \"create_record\", \"args\": {...}}"
    # compact history into text for few-shot context (optional)
    history_text = ""
    for pair in history:
        try:
            u,a = pair[0], pair[1]
            history_text += f"User: {u}\nAssistant: {a}\n"
        except Exception:
            continue
    prompt = f"{system_prompt}\n{history_text}\nUser: {message}\nAssistant:"
    gen = deepseek_generate(prompt, max_tokens=256)
    text = gen.get("text", "")
    # if text looks like JSON with a tool action, try to invoke
    payload = text.strip()
    if payload.startswith("{") or payload.startswith("["):
        try:
            parsed = json.loads(payload)
            if isinstance(parsed, dict) and "tool" in parsed:
                tool_name = parsed.get("tool"); args = parsed.get("args", {})
                if tool_name in globals() and callable(globals()[tool_name]):
                    try:
                        out = globals()[tool_name](**args) if isinstance(args, dict) else globals()[tool_name](args)
                        return f"Invoked tool '{tool_name}'. Result:\n{out}"
                    except Exception as e:
                        return f"Tool invocation error: {e}"
                else:
                    return f"Requested tool '{tool_name}' not found locally."
        except Exception:
            pass
    return text

# ----------------------------
# Gradio chat handler
# ----------------------------
def chat_handler(message, history):
    history = history or []
    trimmed = (message or "").strip()

    # explicit command parser
    cmd = try_parse_and_invoke_command(trimmed)
    if cmd is not None:
        return cmd

    # developer dev path handling (send path unchanged)
    if trimmed.startswith("/mnt/data/"):
        try:
            doc = process_document(trimmed)
            return f"Processed file {doc.get('file')}. Extracted: {json.dumps(doc.get('extracted_data'), ensure_ascii=False)}"
        except Exception as e:
            return f"Error processing document: {e}"

    # otherwise, call deepseek_response (LLM or fallback)
    try:
        return deepseek_response(trimmed, history)
    except Exception as e:
        logger.error("deepseek_response error: %s", e)
        traceback.print_exc()
        return rule_based_response(trimmed)

# ----------------------------
# Gradio UI
# ----------------------------
def chat_interface():
    return gr.ChatInterface(fn=chat_handler, textbox=gr.Textbox(placeholder="Ask me to create contacts, invoices, upload docs (or paste /mnt/data/... for dev)."))

# ----------------------------
# Entrypoint
# ----------------------------
if __name__ == "__main__":
    logger.info("Starting MCP server (DeepSeek mode). Loaded model: %s", LOADED_MODEL_NAME)
    demo = chat_interface()
    demo.launch(server_name="0.0.0.0", server_port=7860)