#!/usr/bin/env python3
from __future__ import annotations
import builtins
import csv, hashlib, hmac, ipaddress, json, os, secrets, threading, time
from collections import defaultdict, deque
from pathlib import Path
from typing import Any, Literal

import mysql.connector
from mysql.connector.pooling import MySQLConnectionPool
from fastapi import BackgroundTasks, Depends, FastAPI, Header, HTTPException, Query, Request
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field

DB_HOST=os.getenv('PD_DB_HOST','127.0.0.1'); DB_PORT=int(os.getenv('PD_DB_PORT','3306'))
DB_USER=os.getenv('PD_DB_USER','directory_api'); DB_PASSWORD=os.getenv('PD_DB_PASSWORD','')
DB_NAME=os.getenv('PD_DB_NAME','public_directory'); SIGN=os.getenv('PD_SIGNING_SECRET','')
EXPORT_DIR=Path(os.getenv('PD_EXPORT_DIR','/home/mostincr/maillists.mostincrediblevalue.com/data/public-directory/exports'))
BASE_URL=os.getenv('PD_BASE_URL','http://127.0.0.1:8086').rstrip('/')
TRUST_PROXY=os.getenv('PD_TRUST_PROXY_HEADERS','false').lower()=='true'
MAX_PAGE=int(os.getenv('PD_MAX_PAGE_SIZE','100')); RATE=int(os.getenv('PD_RATE_LIMIT_PER_MINUTE','120'))
if not DB_PASSWORD: raise RuntimeError('PD_DB_PASSWORD is required')
if len(SIGN)<32: raise RuntimeError('PD_SIGNING_SECRET must be at least 32 characters')
EXPORT_DIR.mkdir(parents=True,exist_ok=True)
pool=MySQLConnectionPool(pool_name='pd',pool_size=int(os.getenv('PD_DB_POOL_SIZE','10')),host=DB_HOST,port=DB_PORT,user=DB_USER,password=DB_PASSWORD,database=DB_NAME,charset='utf8mb4',autocommit=True)
app=FastAPI(title='Public Directory Search API',version='1.0.0')

COLS=['first_name','middle_name','last_name','full_name','title','street_number','street','unit','city','state','zip_code','county','country']
SELECT=', '.join('`'+c+'`' for c in COLS)
SORTS={'relevance':None,'last_name_asc':'`last_name` ASC, `first_name` ASC, `id` ASC','last_name_desc':'`last_name` DESC, `id` DESC','city_asc':'`city` ASC, `last_name` ASC, `id` ASC','zip_asc':'`zip_code` ASC, `last_name` ASC, `id` ASC','street_asc':'`street` ASC, `street_number` ASC, `id` ASC'}
rate_windows=defaultdict(deque); rate_lock=threading.Lock()

def db():
    c=pool.get_connection()
    try: yield c
    finally: c.close()

def cip(req):
    if TRUST_PROXY and req.headers.get('x-forwarded-for'): return req.headers['x-forwarded-for'].split(',')[0].strip()
    return req.client.host if req.client else 'unknown'

def ip_ok(ip, allowed):
    if not allowed: return True
    try: addr=ipaddress.ip_address(ip)
    except ValueError: return False
    for raw in allowed.split(','):
        raw=raw.strip()
        if not raw: continue
        try:
            if ('/' in raw and addr in ipaddress.ip_network(raw,strict=False)) or ('/' not in raw and addr==ipaddress.ip_address(raw)): return True
        except ValueError: pass
    return False

def auth(req:Request,x_api_key:str|None=Header(default=None,alias='X-API-Key'),conn=Depends(db)):
    if not x_api_key: raise HTTPException(401,'Missing API key')
    digest=hashlib.sha256(x_api_key.encode()).hexdigest(); cur=conn.cursor(dictionary=True)
    cur.execute('SELECT id,name,is_active,allowed_ips,rate_limit_per_minute,max_export_records FROM api_keys WHERE key_hash=%s LIMIT 1',(digest,)); row=cur.fetchone(); cur.close()
    if not row or not row['is_active']: raise HTTPException(401,'Invalid API key')
    if not ip_ok(cip(req),row['allowed_ips']): raise HTTPException(403,'IP address not authorized')
    now=time.time(); w=rate_windows[str(row['id'])]; limit=row['rate_limit_per_minute'] or RATE
    with rate_lock:
        while w and w[0]<now-60: w.popleft()
        if len(w)>=limit: raise HTTPException(429,'Rate limit exceeded')
        w.append(now)
    return row

def clean(v): return ' '.join(v.strip().split()) if isinstance(v,str) and v.strip() else None

def bq(q):
    toks=[]
    for raw in q.split():
        t=''.join(ch for ch in raw if ch.isalnum() or ch in "-'")
        if len(t)>=2: toks.append('+'+t+'*')
    return ' '.join(toks)

class F(BaseModel):
    q:str|None=Field(None,max_length=250); first_name:str|None=None; middle_name:str|None=None; last_name:str|None=None
    street_number:str|None=None; street:str|None=None; city:str|None=None; state:str|None=None; zip:str|None=None; county:str|None=None
    first_name_match:Literal['exact','starts_with','contains']='starts_with'; last_name_match:Literal['exact','starts_with','contains']='starts_with'
    street_match:Literal['exact','starts_with','contains']='contains'; city_match:Literal['exact','starts_with','contains']='starts_with'
class CountReq(BaseModel): filters:F
class ExportReq(BaseModel): order_id:str=Field(min_length=1,max_length=100); search_snapshot:F; record_quantity:int=Field(ge=1,le=1_000_000)

def where(f:F):
    clauses=[]; params=[]; rel=None
    if clean(f.q) and bq(clean(f.q)):
        rel='MATCH(first_name,middle_name,last_name,full_name,street,city,county) AGAINST (%s IN BOOLEAN MODE)'; clauses.append(rel); params.append(bq(clean(f.q)))
    for col,val,mode in [('first_name',f.first_name,f.first_name_match),('middle_name',f.middle_name,'starts_with'),('last_name',f.last_name,f.last_name_match),('street_number',f.street_number,'exact'),('street',f.street,f.street_match),('city',f.city,f.city_match),('state',f.state,'exact'),('county',f.county,'starts_with')]:
        val=clean(val)
        if not val: continue
        if mode=='exact': clauses.append(f'`{col}`=%s'); params.append(val)
        elif mode=='starts_with': clauses.append(f'`{col}` LIKE %s'); params.append(val+'%')
        else: clauses.append(f'`{col}` LIKE %s'); params.append('%'+val+'%')
    z=clean(f.zip)
    if z:
        if len(z)==5 and z.isdigit(): clauses.append('`zip_code` LIKE %s'); params.append(z+'%')
        else: clauses.append('`zip_code`=%s'); params.append(z)
    return ' AND '.join(clauses) if clauses else '1=1',params,rel

def usage(conn,key_id,endpoint,ip,n=0):
    cur=conn.cursor(); cur.execute('INSERT INTO api_usage(api_key_id,endpoint,ip_address,records_returned) VALUES(%s,%s,%s,%s)',(key_id,endpoint,ip,n)); cur.close()

@app.get('/health')
def health(conn=Depends(db)):
    cur=conn.cursor(); cur.execute('SELECT 1'); cur.fetchone(); cur.close(); return {'status':'ok','version':'1.0.0'}

@app.get('/v1/search')
def search(request:Request,q:str|None=None,first_name:str|None=None,middle_name:str|None=None,last_name:str|None=None,street_number:str|None=None,street:str|None=None,city:str|None=None,state:str|None=None,zip:str|None=None,county:str|None=None,first_name_match:Literal['exact','starts_with','contains']='starts_with',last_name_match:Literal['exact','starts_with','contains']='starts_with',street_match:Literal['exact','starts_with','contains']='contains',city_match:Literal['exact','starts_with','contains']='starts_with',page:int=Query(1,ge=1,le=10000),page_size:int=Query(25,ge=1),sort:Literal['relevance','last_name_asc','last_name_desc','city_asc','zip_asc','street_asc']='relevance',key=Depends(auth),conn=Depends(db)):
    page_size=min(page_size,MAX_PAGE); f=F(q=q,first_name=first_name,middle_name=middle_name,last_name=last_name,street_number=street_number,street=street,city=city,state=state,zip=zip,county=county,first_name_match=first_name_match,last_name_match=last_name_match,street_match=street_match,city_match=city_match)
    w,p,rel=where(f); cur=conn.cursor(); cur.execute('SELECT COUNT(*) FROM people_addresses WHERE '+w,tuple(p)); total=int(cur.fetchone()[0])
    select=SELECT; rp=list(p); order=SORTS[sort]
    if sort=='relevance' and rel:
        select+=', '+rel+' AS relevance'; rp=[bq(clean(q) or '')]+rp; order='relevance DESC, `last_name` ASC, `id` ASC'
    elif sort=='relevance': order='`last_name` ASC, `first_name` ASC, `id` ASC'
    rp += [page_size,(page-1)*page_size]; cur.execute(f'SELECT {select} FROM people_addresses WHERE {w} ORDER BY {order} LIMIT %s OFFSET %s',tuple(rp)); names=[d[0] for d in cur.description]; results=[dict(builtins.zip(names,r)) for r in cur.fetchall()]; cur.close(); usage(conn,key['id'],'/v1/search',cip(request),len(results)); return {'total':total,'page':page,'page_size':page_size,'results':results}

@app.post('/v1/search/count')
def count(payload:CountReq,request:Request,key=Depends(auth),conn=Depends(db)):
    w,p,_=where(payload.filters); cur=conn.cursor(); cur.execute('SELECT COUNT(*) FROM people_addresses WHERE '+w,tuple(p)); total=int(cur.fetchone()[0]); cur.close(); usage(conn,key['id'],'/v1/search/count',cip(request)); return {'total':total}

AUTO={'cities':('city',['state']),'zips':('zip_code',['city','state']),'streets':('street',['city','state']),'first-names':('first_name',[]),'last-names':('last_name',[])}
@app.get('/v1/autocomplete/{kind}')
def ac(kind:str,request:Request,q:str=Query(min_length=2,max_length=100),city:str|None=None,state:str|None=None,zip:str|None=None,key=Depends(auth),conn=Depends(db)):
    if kind not in AUTO: raise HTTPException(404,'Unknown autocomplete type')
    col,extras=AUTO[kind]; clauses=[f'`{col}` LIKE %s',f'`{col}` IS NOT NULL',f'`{col}`<>\'\'']; p=[clean(q)+'%']
    if city and col!='city': clauses.append('`city`=%s'); p.append(clean(city))
    if state: clauses.append('`state`=%s'); p.append(clean(state))
    if zip and col!='zip_code': clauses.append('`zip_code` LIKE %s'); p.append(clean(zip)+'%')
    fields=[col]+extras; sel=', '.join('`'+x+'`' for x in fields); grp=sel; cur=conn.cursor(dictionary=True); cur.execute(f'SELECT {sel},COUNT(*) cnt FROM people_addresses WHERE {" AND ".join(clauses)} GROUP BY {grp} ORDER BY cnt DESC, `{col}` ASC LIMIT 10',tuple(p)); rows=cur.fetchall(); cur.close(); out=[]
    for r in rows:
        label=r[col]
        if kind=='cities' and r.get('state'): label=f"{r[col]}, {r['state']}"
        elif kind=='zips': label=f"{r[col]} — {r.get('city','')}, {r.get('state','')}".strip(' ,—')
        out.append({'value':r[col],'label':label,'count':int(r['cnt'])})
    usage(conn,key['id'],'/v1/autocomplete/'+kind,cip(request)); return {'suggestions':out}

def token(exp_id,expires):
    sig=hmac.new(SIGN.encode(),f'{exp_id}:{expires}'.encode(),hashlib.sha256).hexdigest(); return f'{expires}.{sig}'
def verify_token(exp_id,t):
    try: es,sig=t.split('.',1); e=int(es)
    except: raise HTTPException(403,'Invalid download token')
    if e<int(time.time()): raise HTTPException(403,'Download link expired')
    ex=hmac.new(SIGN.encode(),f'{exp_id}:{e}'.encode(),hashlib.sha256).hexdigest()
    if not hmac.compare_digest(sig,ex): raise HTTPException(403,'Invalid download token')

def do_export(exp_id,fdict,qty):
    conn=pool.get_connection(); cur=None
    try:
        f=F(**fdict); w,p,rel=where(f); cur=conn.cursor(); cur.execute("UPDATE export_jobs SET status='generating',started_at=NOW() WHERE export_id=%s",(exp_id,))
        if rel and f.q:
            sql=f'SELECT {SELECT}, {rel} AS relevance FROM people_addresses WHERE {w} ORDER BY relevance DESC,`last_name` ASC,`id` ASC LIMIT %s'; qp=[bq(clean(f.q) or '')]+p+[qty]
        else:
            sql=f'SELECT {SELECT} FROM people_addresses WHERE {w} ORDER BY `last_name` ASC,`first_name` ASC,`id` ASC LIMIT %s'; qp=p+[qty]
        cur.execute(sql,tuple(qp)); path=EXPORT_DIR/f'{exp_id}.csv'; written=0
        with path.open('w',encoding='utf-8',newline='') as fh:
            wr=csv.writer(fh); wr.writerow(['First Name','Middle Name','Last Name','Full Name','Title','Street Number','Street','Unit','City','State','ZIP Code','County','Country'])
            while True:
                rows=cur.fetchmany(5000)
                if not rows: break
                for row in rows: wr.writerow(row[:len(COLS)]); written+=1
        cur.execute("UPDATE export_jobs SET status='ready',record_count=%s,file_path=%s,ready_at=NOW(),last_error=NULL WHERE export_id=%s",(written,str(path),exp_id))
    except Exception as e:
        try:
            c=conn.cursor(); c.execute("UPDATE export_jobs SET status='failed',last_error=%s WHERE export_id=%s",(str(e)[:1000],exp_id)); c.close()
        except: pass
    finally:
        if cur: cur.close()
        conn.close()

@app.post('/v1/exports')
def create_export(payload:ExportReq,bg:BackgroundTasks,request:Request,key=Depends(auth),conn=Depends(db)):
    maxr=key['max_export_records'] or 100000
    if payload.record_quantity>maxr: raise HTTPException(403,f'Export exceeds key limit of {maxr:,} records')
    w,p,_=where(payload.search_snapshot); cur=conn.cursor(); cur.execute('SELECT COUNT(*) FROM people_addresses WHERE '+w,tuple(p)); total=int(cur.fetchone()[0])
    if payload.record_quantity>total: cur.close(); raise HTTPException(400,f'Only {total:,} records match this search')
    cur.execute('SELECT export_id,status FROM export_jobs WHERE order_id=%s AND api_key_id=%s LIMIT 1',(payload.order_id,key['id'])); old=cur.fetchone()
    if old: cur.close(); return {'export_id':old[0],'status':old[1]}
    eid='exp_'+secrets.token_urlsafe(18); cur.execute("INSERT INTO export_jobs(export_id,order_id,api_key_id,status,requested_count,search_snapshot) VALUES(%s,%s,%s,'queued',%s,%s)",(eid,payload.order_id,key['id'],payload.record_quantity,json.dumps(payload.search_snapshot.model_dump(),separators=(',',':')))); cur.close(); bg.add_task(do_export,eid,payload.search_snapshot.model_dump(),payload.record_quantity); usage(conn,key['id'],'/v1/exports',cip(request)); return {'export_id':eid,'status':'queued'}

@app.get('/v1/exports/{eid}')
def status(eid:str,request:Request,key=Depends(auth),conn=Depends(db)):
    cur=conn.cursor(dictionary=True); cur.execute('SELECT export_id,status,requested_count,record_count,created_at,ready_at,last_error FROM export_jobs WHERE export_id=%s AND api_key_id=%s LIMIT 1',(eid,key['id'])); r=cur.fetchone(); cur.close()
    if not r: raise HTTPException(404,'Export not found')
    if r['status']=='ready':
        e=int(time.time())+3600; r['download_url']=f'{BASE_URL}/v1/exports/{eid}/download?token={token(eid,e)}'
    else: r['download_url']=None
    for k in ('created_at','ready_at'):
        if r.get(k): r[k]=r[k].isoformat()
    usage(conn,key['id'],'/v1/exports/'+eid,cip(request)); return r

@app.get('/v1/exports/{eid}/download')
def download(eid:str,token:str,conn=Depends(db)):
    verify_token(eid,token); cur=conn.cursor(dictionary=True); cur.execute('SELECT status,file_path FROM export_jobs WHERE export_id=%s LIMIT 1',(eid,)); r=cur.fetchone()
    if not r or r['status']!='ready' or not r['file_path']: cur.close(); raise HTTPException(404,'Export not ready')
    p=Path(r['file_path'])
    if not p.is_file() or p.parent.resolve()!=EXPORT_DIR.resolve(): cur.close(); raise HTTPException(404,'Export file missing')
    cur.execute('UPDATE export_jobs SET download_count=download_count+1,first_download_at=COALESCE(first_download_at,NOW()),last_download_at=NOW() WHERE export_id=%s',(eid,)); cur.close(); return FileResponse(str(p),media_type='text/csv',filename=f'mailing-list-{eid}.csv')