import os
import logging
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from langchain_ollama import OllamaLLM
from langchain_core.callbacks.manager import CallbackManager
from langchain_core.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from jira_mcp_connector import JiraMCPConnector
from telegramchat import TelegramChatService

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="Qwen2.5 7B on HF Space")

# ================== CHANGE HERE ==================
MODEL_NAME = "qwen2.5:7b-instruct"   # Best chat version
# ================================================

def get_llm():
    callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
    return OllamaLLM(
        model=MODEL_NAME,
        temperature=0.7,
        num_ctx=8192,        # Qwen supports bigger context
        callback_manager=callback_manager,
    )

class Question(BaseModel):
    text: str


jira_connector = JiraMCPConnector(
    llm_factory=get_llm,
    logger=logger,
)


telegram_chat_service = TelegramChatService(
    llm_factory=get_llm,
    token=os.getenv("TELEGRAM_TOKEN", ""),
    logger=logger,
    jira_connector=jira_connector,
)

@app.get("/")
def read_root():
    return {
        "message": "✅ Qwen2.5 7B is running!",
        "model": MODEL_NAME,
        "endpoints": ["/ask", "/stream"]
    }

@app.post("/ask")
async def ask(question: Question):
    try:
        logger.info(f"Question: {question.text[:100]}...")
        llm = get_llm()
        response = llm.invoke(question.text)
        return {
            "question": question.text,
            "answer": response.strip()
        }
    except Exception as e:
        logger.error(f"Error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/stream")
async def stream():
    try:
        jira_connector.ensure_services_started()
        started = telegram_chat_service.start()
        return {
            "telegram_running": telegram_chat_service.is_running,
            "jira_api_running": True,
            "status": "started" if started else "already_running",
            "message": "Telegram + Jira MCP flow is active. Open Telegram and message your bot.",
        }
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"Error starting Telegram chat: {e}")
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))