from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import pipeline
import re

app = FastAPI(
    title="AI Text Cleaner API",
    version="1.0"
)

# CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load model
pipe = pipeline(
    "text-generation",
    model="HuggingFaceTB/SmolLM2-1.7B-Instruct",
    device_map="cpu"
)

# Request Model
class TextRequest(BaseModel):
    text: str
    mode: str = "clean"

# Fast local cleanup before AI
def basic_cleanup(text):
    
    # remove extra spaces
    text = re.sub(r'\s+', ' ', text)

    # remove repeated empty lines
    text = re.sub(r'\n+', '\n', text)

    # trim
    text = text.strip()

    return text

# Prompt templates
PROMPTS = {

    "clean": """
You are an AI text cleaning assistant.

Tasks:
- Remove messy formatting
- Fix weird spacing
- Improve readability
- Keep original meaning
- Return only cleaned text

Text:
{text}
""",

    "ocr": """
You are an OCR text repair assistant.

Tasks:
- Fix OCR mistakes
- Merge broken words
- Remove strange symbols
- Fix broken lines
- Return only repaired text

Text:
{text}
""",

    "format": """
You are a text formatting assistant.

Tasks:
- Format text properly
- Improve structure
- Fix capitalization
- Make text readable
- Return only formatted text

Text:
{text}
""",

    "prompt": """
You are an AI prompt optimizer.

Tasks:
- Clean the prompt
- Improve clarity
- Keep original intent
- Make prompt concise
- Return only optimized prompt

Prompt:
{text}
"""
}

@app.get("/")
async def root():
    return {
        "status": "running",
        "service": "AI Text Cleaner API"
    }

@app.post("/process")
async def process_text(req: TextRequest):

    # input limit
    if len(req.text) > 3000:
        return {
            "error": "Text too long"
        }

    # local cleanup
    cleaned_input = basic_cleanup(req.text)

    # select prompt
    prompt_template = PROMPTS.get(req.mode, PROMPTS["clean"])

    prompt = prompt_template.format(text=cleaned_input)

    # generate
    result = pipe(
        prompt,
        max_new_tokens=180,
        temperature=0.2,
        do_sample=False
    )

    output = result[0]["generated_text"]

    # remove prompt echo
    if prompt in output:
        output = output.replace(prompt, "").strip()

    return {
        "success": True,
        "mode": req.mode,
        "result": output
    }