from fastapi import FastAPI
from pydantic import BaseModel
from typing import List
from src.predictive.custom_model import ModelAPI
from src.slm.core import configure_llm
from src.slm.operation import workflow
from src.slm.query import xia_query, summary_query
from src.slm.retrieval import Retrieval

modelAPI = ModelAPI()
configure_llm()
agent = workflow()
retrieval = Retrieval()

app = FastAPI(title="Recon Robot Health Agent", description="Recon Robot Health Agent API")

class SensorInput(BaseModel):
    instances: List[float]

class QueryInput(BaseModel):
    query: str

def alert_message(status):
  return f"The status of the thruster is {status}"

async def identify_root_cause(agent, thruster_data):
  sensor_data = [{**thruster_data}]
  factor_query = xia_query(sensor_data[0])
  sum_query = summary_query()
  return await agent.run(user_msg=f"{factor_query} {sum_query}")

async def search_document(retrieval, text):
  return await retrieval.query_context(text)


@app.post("/predict")
def predict(data: SensorInput):
    features = {
        'voltage': data.instances[0],
        'current': data.instances[1],
        'power': data.instances[2],
        'temperature': data.instances[3],
        'driver_temperature': data.instances[4],
        'speed_rpm': data.instances[5],
        'thruster_id_encoded': data.instances[6]
        }
    prediction = modelAPI.predict(features)
    status = alert_message(prediction)

    return {"prediction": status}


@app.post("/rca")
async def rca(data: SensorInput):
    features = {
        'voltage': data.instances[0],
        'current': data.instances[1],
        'power': data.instances[2],
        'temperature': data.instances[3],
        'driver_temperature': data.instances[4],
        'speed_rpm': data.instances[5],
        'thruster_id_encoded': data.instances[6]
        }
    long_response = await identify_root_cause(agent,  features)

    return {"rca": str(long_response)}


@app.post('/semantic')
async def semantic(data: QueryInput):
  response =  await search_document(retrieval, data.query)
  return {'search': str(response)}