import io
from typing import Optional
from fastapi import FastAPI, File, UploadFile, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from pypdf import PdfReader
import pytesseract
from PIL import Image

app = FastAPI()

# CORS configuration
origins = [
    "http://localhost",
    "http://localhost:3000",
    "http://localhost:5173", # Vite default
    "http://localhost:4200", # Angular default
    "http://127.0.0.1:4200",
]

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

# Models
class LoginRequest(BaseModel):
    username: str
    password: str

class LoginResponse(BaseModel):
    token: str
    message: str

class NERResponse(BaseModel):
    text: str

# 1. User Login Module (Dummy)
@app.post("/api/login", response_model=LoginResponse)
async def login(request: LoginRequest):
    if request.username == "admin" and request.password == "admin":
        return LoginResponse(token="dummy-jwt-token-123", message="Login Successful")
    raise HTTPException(status_code=401, detail="Invalid credentials")

# 2. OCR Module
def extract_text_from_pdf(file_bytes: bytes) -> str:
    try:
        reader = PdfReader(io.BytesIO(file_bytes))
        text = ""
        for page in reader.pages:
            page_text = page.extract_text()
            if page_text:
                text += page_text + "\n"
        return text.strip()
    except Exception as e:
        print(f"Error reading PDF: {e}")
        return ""

def extract_text_from_image(file_bytes: bytes) -> str:
    try:
        image = Image.open(io.BytesIO(file_bytes))
        text = pytesseract.image_to_string(image)
        return text.strip()
    except Exception as e:
        print(f"Error reading Image: {e}")
        return ""

from pdf2image import convert_from_bytes

# ... (imports)

@app.post("/api/ocr/extract", response_model=NERResponse)
async def extract_text(file: UploadFile = File(...)):
    content = await file.read()
    filename = file.filename.lower()
    
    extracted_text = ""
    
    if filename.endswith(".pdf"):
        # Try text extraction first
        extracted_text = extract_text_from_pdf(content)
        
        # If text is empty, it might be a scanned PDF. 
        if not extracted_text.strip():
             try:
                 images = convert_from_bytes(content)
                 for i, image in enumerate(images):
                     page_text = pytesseract.image_to_string(image)
                     extracted_text += f"\n--- Page {i+1} ---\n{page_text}"
             except Exception as e:
                 extracted_text = f"Error processing Scanned PDF: {str(e)}\n\n(Hint: Ensure 'poppler' is installed on your system. Run 'brew install poppler')"

    elif filename.endswith((".png", ".jpg", ".jpeg", ".tiff", ".bmp")):
        extracted_text = extract_text_from_image(content)
    else:
        raise HTTPException(status_code=400, detail="Unsupported file type")
    
    return NERResponse(text=extracted_text)

@app.get("/")
def read_root():
    return {"message": "OCR Backend API is running"}