"""Validate queries on disposable native MySQL 8/PostgreSQL 16 temporary tables."""
import argparse
import json
import re
import time
from pathlib import Path
import sqlglot
from sqlglot import exp
from tinyquery.data import ddls,fixture,SQL_OPS


class NativeDB:
    def __init__(self,backend):
        self.backend=backend
        if backend=='mysql':
            import pymysql
            self.conn=pymysql.connect(unix_socket='/run/mysqld/mysqld.sock',user='root',database='tinyquery',autocommit=True)
        else:
            import psycopg
            self.conn=psycopg.connect(host='/var/run/postgresql',user='root',dbname='postgres',autocommit=True)
        self.cursor=self.conn.cursor(); self.tables=[]
        self.cursor.execute('SET SESSION max_execution_time=2000' if backend=='mysql' else "SET statement_timeout='2s'")
    def setup(self,s,seed):
        for table in reversed(self.tables): self.cursor.execute(f'DROP TABLE IF EXISTS {table}')
        self.tables=[]
        for ddl in ddls(s):
            # Relationships are populated identically; MySQL temp tables cannot enforce foreign keys.
            ddl=re.sub(r' REFERENCES [a-zA-Z_]+\(id\)','',ddl)
            self.cursor.execute(ddl.replace('CREATE TABLE','CREATE TEMPORARY TABLE',1))
        self.tables=[s['parent'],s['table']]
        source=fixture(s,seed)
        for table in self.tables:
            rows=source.execute(f'SELECT * FROM {table}').fetchall()
            self.cursor.executemany(f"INSERT INTO {table} VALUES ({','.join(['%s']*len(rows[0]))})",rows)
        source.close()
    def execute(self,sql):
        statements=sqlglot.parse(sql,read='mysql' if self.backend=='mysql' else 'postgres')
        if len(statements)!=1 or not isinstance(statements[0],exp.Select): raise ValueError('Expected one SELECT')
        tree=statements[0]
        if tree.find(exp.Into): raise ValueError('SELECT INTO is not permitted')
        if any(isinstance(node,(exp.DML,exp.DDL)) for node in tree.walk()):
            raise ValueError('Data-changing statements are not permitted inside a query')
        if any(t.name not in self.tables for t in tree.find_all(exp.Table)): raise ValueError('Query references an unknown table')
        for f in tree.find_all(exp.Anonymous):
            if f.name.upper() not in ['YEAR','MONTH','LOWER','UPPER']:
                raise ValueError('Unsupported function '+f.name)
        self.cursor.execute(sql)
        rows=self.cursor.fetchmany(1001)
        if len(rows)>1000: raise ValueError('Result limit exceeded')
        return [tuple(str(x) if not isinstance(x,(str,int,float,type(None))) else x for x in row) for row in rows]
    def close(self): self.cursor.close(); self.conn.close()


def main():
    p=argparse.ArgumentParser(); p.add_argument('--data',required=True); p.add_argument('--predictions')
    p.add_argument('--out',required=True); p.add_argument('--limit',type=int,default=0)
    args=p.parse_args()
    rows=[json.loads(l) for l in Path(args.data).read_text().splitlines()]
    predictions={}
    if args.predictions: predictions={r['id']:r for r in map(json.loads,Path(args.predictions).read_text().splitlines())}
    dbs={b:NativeDB(b) for b in ['mysql','supabase']}; results=[]; seen=set(); start=time.time()
    for r in rows:
        if r['operation'] not in SQL_OPS: continue
        if not args.predictions:
            key=(r['operation'],r['backend'])
            if key in seen: continue
            seen.add(key)
        elif r['id'] not in predictions: continue
        if args.limit and len(results)>=args.limit: break
        db=dbs[r['backend']]
        item={'id':r['id'],'backend':r['backend'],'operation':r['operation'],'gold_executes':False,'prediction_equivalent':None}
        try:
            gold_key=next(k for k in ('query','sql') if k in r['target']['arguments'])
            gold=r['target']['arguments'][gold_key]
            prediction=None
            if args.predictions:
                item['prediction_equivalent']=False
                from tinyquery.evaluate import check_action,parse_action
                action=parse_action(predictions[r['id']]['output'])
                check_action(action,r['context'])
                if action.get('name')!=r['target']['name']: raise ValueError('Wrong tool')
                expected_other={k:v for k,v in r['target']['arguments'].items() if k!=gold_key}
                actual_other={k:v for k,v in action['arguments'].items() if k!=gold_key}
                if expected_other!=actual_other: raise ValueError('Wrong non-SQL arguments')
                prediction=action['arguments'][gold_key]
            equal=True
            for seed in [1,7]:
                db.setup(r['slots'],seed); expected=db.execute(gold)
                if prediction is not None:
                    actual=db.execute(prediction)
                    ordered=r['operation'] in ['sort_asc','sort_desc','top','bottom']
                    if not ordered: expected=sorted(expected,key=repr); actual=sorted(actual,key=repr)
                    equal=equal and actual==expected
            item['gold_executes']=True
            if prediction is not None: item['prediction_equivalent']=equal
        except Exception as exc: item['error']=str(exc)[:300]
        results.append(item)
    for db in dbs.values(): db.close()
    summary={'seconds':time.time()-start,'cases':len(results),'gold_executes':sum(r['gold_executes'] for r in results),
             'prediction_equivalent':sum(r['prediction_equivalent'] is True for r in results) if args.predictions else None,
             'engine_versions':{'mysql':'8.0.46','postgresql':'16.15'},'results':results}
    Path(args.out).write_text(json.dumps(summary,indent=2))
    print(json.dumps({k:v for k,v in summary.items() if k!='results'}),flush=True)
    for r in results:
        if 'error' in r: print(json.dumps(r),flush=True)


if __name__=='__main__': main()