import dspy
import os

from dspy.predict.react import Tool
from tavily import TavilyClient


#lm = dspy.LM('ollama_chat/llama3.2', api_base='http://localhost:11434', api_key='')
#lm = dspy.LM('ollama_chat/deepseek-r1', api_base='http://localhost:11434', api_key='')
#lm = dspy.LM('ollama_chat/qwq', api_base='http://localhost:11434', api_key='')
#lm = dspy.LM('ollama_chat/deepscaler', api_base='http://localhost:11434', api_key='')
#lm = dspy.LM('huggingface/Qwen/Qwen2.5-Coder-32B-Instruct')
#lm = dspy.LM('huggingface/meta-llama/Llama-3.2-1B')
lm = dspy.LM('groq/llama-3.1-8b-instant', api_base='https://api.groq.com/openai/v1', api_key=os.environ["GROQ_API_KEY"])
dspy.configure(lm=lm)

search_client = TavilyClient(api_key=os.environ["T_TOKEN"])

class RedFlagsGenerator(dspy.Signature):
  """Generate red flags based on story in fraud topologies"""
  story: str = dspy.InputField()
  red_flag: str = dspy.OutputField()

class ControlsGenerator(dspy.Signature):
  """Generate preventive, detective, corrective controls based on red flag"""
  red_flag: str = dspy.InputField()
  control: str = dspy.OutputField()

def search_wikipedia(query: str) -> list[str]:
  """Query ColBERT endpoint, which is a knowledge source based on wikipedia data"""
  results = dspy.ColBERTv2(url='http://20.102.90.50:2017/wiki17_abstracts')(query, k=1)
  return [x["text"] for x in results]


def control_search(query: str) -> list[str]:
  """Run a web search to return the top 3 control to verify red flags is a hoax or false positive"""
  response = search_client.search(query)
  return [r["content"] for r in response["results"]]

class Detective(dspy.Module):
  """
  As a fraud investigator, analyze the following case story to assess the likelihood of fraud. Use a structured framework that includes:

  Fraudster Profile & Actions
  Identify the behavioral profile of the fraudster (e.g., boaster, manipulator, deceiver).
  Map their actions to common fraud typologies (social engineering, impersonation, identity theft, financial asset targeting).

  Highlight red flags (fake phone numbers, urgent secrecy, impersonation of authorities, threats of jail).
  Use control to validate whether each red flag is hoax or a false positive.

  Examples: cross-checking official contact numbers on verified government websites, confirming with independent third parties, requesting written documentation, checking with financial institutions’ fraud units, or using telecom records (Rogers, etc.) for number spoofing.
  """
  def __init__(self):
    self.redFlags_generator = dspy.ChainOfThought(RedFlagsGenerator)
    self.controls_generator = dspy.ReAct(ControlsGenerator, tools=[Tool(control_search), Tool(search_wikipedia)])

  def forward(self, story, **kwargs):
    fraud = self.redFlags_generator(story=story).red_flag
    return self.controls_generator(red_flag=fraud).control



story = """
government SCAM PHONE CALL-> identity stolen(provide time and location, force you to contact police(scammer actor)) -> cooperate(social engineering) -> (make you involved in confidential case, dont tell anyone or in jail for 6 month, good cop bad cop) you are suspect -> bail to secure account
"""

conan = Detective()
#dspy.inspect_history()

if __name__=='__main__':
  #recommend(story)

  print(conan(story=story))