from fastapi import FastAPI
from pydantic import BaseModel
from fastapi.responses import JSONResponse
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from openai import OpenAI

print("Version ---- DeepSeek Only")
app = FastAPI()

# -----------------------------
# Request schema
# -----------------------------
class ConflictDetectionRequest(BaseModel):
    Req1: str
    Req2: str
    model_choice: str   # "GPT-4", "DeepSeek-Reasoner", "Fanar"
    prompt_type: str    # "zero-shot" or "few-shot"
    api_key: str = None # required only if model_choice == "GPT-4"

# -----------------------------
# Prompt builder
# -----------------------------
def build_prompt(req1, req2, prompt_type="zero-shot"):
    if prompt_type == "zero-shot":
        return f"Do the following sentences contradict each other, Answer with only 'yes' or 'no', no explanation. \n 1.{req1} 2.{req2}"
    elif prompt_type == "few-shot":
        examples = (
            "Example 1:\n"
            "Req1: The system shall allow password reset.\n"
            "Req2: The system shall not allow password reset.\n"
            "Answer: yes\n\n"
            "Example 2:\n"
            "Req1: The system shall support Arabic language.\n"
            "Req2: The system shall support English language.\n"
            "Answer: no\n\n"
        )
        return examples + f"Now answer: Do the following sentences contradict each other? 1.{req1} 2.{req2}"
    else:
        return f"Do the following sentences contradict each other, yes or no: 1.{req1} 2.{req2}"

# -----------------------------
# Startup: load DeepSeek once
# -----------------------------
@app.on_event("startup")
def load_models():
    print("Loading DeepSeek model into memory...")
    model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
    app.state.deepseek_tokenizer = AutoTokenizer.from_pretrained(model_name)
    app.state.deepseek_tokenizer.pad_token = app.state.deepseek_tokenizer.eos_token
    
    app.state.deepseek_model = AutoModelForCausalLM.from_pretrained(
        model_name,
        # max_tokens=10,  # Limit to very few tokens
        temperature=0.1,    # Low temperature for more deterministic output
        do_sample=False,    # Disable sampling for consistent results
        torch_dtype="auto", 
        device_map="auto", 
        offload_folder="offload"
    )


# -----------------------------
# Model handlers
# -----------------------------
def run_gpt4(req1, req2, prompt_type, api_key):
    client = OpenAI(base_url="https://openrouter.ai/api/v1", api_key=api_key)
    prompt = build_prompt(req1, req2, prompt_type)
    completion = client.chat.completions.create(
        model="openai/gpt-4",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.7,
        max_tokens=512
    )
    return completion.choices[0].message.content.strip()


def run_deepseek(req1, req2, prompt_type):
    tokenizer = app.state.deepseek_tokenizer
    model = app.state.deepseek_model
    prompt = build_prompt(req1, req2, prompt_type)
    inputs = tokenizer([prompt], return_tensors="pt", padding=True, truncation=True)
    outputs = model.generate(
        input_ids=inputs.input_ids,
        attention_mask=inputs.attention_mask,
        max_new_tokens=256,
        pad_token_id=tokenizer.eos_token_id
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def run_fanar(req1, req2, prompt_type):
    client = OpenAI(base_url="https://api.fanar.qa/v1", api_key=os.getenv("FANAR_API"))
    prompt = build_prompt(req1, req2, prompt_type)
    response = client.chat.completions.create(
        model="Fanar",
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content.strip()

# -----------------------------
# API route
# -----------------------------
@app.post("/predict")
def predict(request: ConflictDetectionRequest):
    try:
        if request.model_choice == "GPT-4":
            if not request.api_key:
                return JSONResponse({"error": "API key required for GPT-4"}, status_code=400)
            answer = run_gpt4(request.Req1, request.Req2, request.prompt_type, request.api_key)

        elif request.model_choice == "DeepSeek-Reasoner":
            answer = run_deepseek(request.Req1, request.Req2, request.prompt_type)

        elif request.model_choice == "Fanar":
            answer = run_fanar(request.Req1, request.Req2, request.prompt_type)

        else:
            return JSONResponse({"error": "Invalid model_choice"}, status_code=400)

        return JSONResponse({"resp": answer, "statusText": "OK", "statusCode": 0}, status_code=200)

    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)


def run_fanar_to_resolve(req1, req2):
    client = OpenAI(base_url="https://api.fanar.qa/v1", api_key=os.getenv("FANAR_API"))
    prompt = f"""
You are a requirements engineer. You are given two potentially conflicting requirements.


Requirement 1: "{req1}"
Requirement 2: "{req2}"


Your task is to resolve the conflict in Arabic using a hybrid strategy:
1. If possible, combine both requirements into **a single non-conflicting and concise requirement** that preserves the intent of both.
2. If they cannot be merged clearly, rewrite them as **two separate non-conflicting requirements**, possibly by adding conditions, default cases, or clarifying their scope.


Avoid using vague expressions such as 'as needed', 'if possible', or 'may'. Only merge requirements if they can be resolved clearly and deterministically. Otherwise, indicate that clarification is needed.


The resolved output must follow proper software requirement language and be easy to test.


Format your answer like this:
[Resolution Type]: Single / Split
Resolved Requirement(s):
- ...
"""
    
    response = client.chat.completions.create(
        model="Fanar",
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content.strip()


class ConflictResoluionRequest(BaseModel):
    Req1: str
    Req2: str


@app.post("/resolve")
def resolve(request: ConflictResoluionRequest):
    answer = run_fanar_to_resolve(request.Req1, request.Req2)
    return JSONResponse({"resp": answer, "statusText": "OK", "statusCode": 0}, status_code=200)