#include "../include/slime/upper_interval.hpp"
#include <cuda_runtime.h>
#include <algorithm>
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cmath>
#include <cstdint>
#include <deque>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <mutex>
#include <set>
#include <sstream>
#include <stdexcept>
#include <string>
#include <thread>
#include <vector>
namespace sr=slime_research;
namespace fs=full_search_experiment;
using Clock=std::chrono::steady_clock;
#define CUDA(call) do{cudaError_t status_=(call);if(status_!=cudaSuccess)throw std::runtime_error(std::string(#call)+": "+cudaGetErrorString(status_));}while(0)
struct Coord{int32_t x,z;};
struct Job{uint32_t position,root;};
struct Match{uint64_t seed;int32_t x,z;};
struct Config{int32_t x=0,z=0;uint64_t radius=0;int size=5,positionBatch=32,jobBatch=2048,repeats=1;uint64_t maxResults=0;size_t cacheClasses=4096;std::string devices="0",output;bool dry=false,list=false;};
uint64_t number(const std::string& s){size_t used=0;if(s.empty()||s[0]=='-')throw std::runtime_error("expected nonnegative integer: "+s);auto n=std::stoull(s,&used);if(used!=s.size())throw std::runtime_error("invalid integer: "+s);return n;}
int32_t signed32(const std::string&s){size_t used=0;auto n=std::stoll(s,&used);if(used!=s.size()||n<INT32_MIN||n>INT32_MAX)throw std::runtime_error("invalid int32 coordinate");return int32_t(n);}
uint64_t anchors(uint64_t r){return 1+2*r*(r+1);}
uint64_t isqrt(uint64_t v){uint64_t r=std::sqrt(double(v));while((r+1)*(r+1)<=v)++r;while(r*r>v)--r;return r;}
Coord coordinate(const Config& c,uint64_t index){
    uint64_t top=(c.radius+1)*(c.radius+1);int64_t dx,dz;
    if(index<top){uint64_t row=isqrt(index);dz=int64_t(row)-c.radius;dx=int64_t(index-row*row)-int64_t(row);}
    else{uint64_t back=anchors(c.radius)-1-index,row=isqrt(back);dz=int64_t(c.radius)-row;dx=int64_t(row)-int64_t(back-row*row);}
    return {int32_t(int64_t(c.x)+dx),int32_t(int64_t(c.z)+dz)};
}
uint64_t classKey(Coord p){return (uint64_t(uint32_t(p.x)&65535)<<17)|(uint32_t(p.z)&131071);}
void terms(Coord p,int n,sr::CellTerm* out){uint64_t anchor=sr::coordinateTerm(p.x,p.z);for(int z=0;z<n;++z)for(int x=0;x<n;++x)out[z*n+x]={(sr::coordinateTerm(p.x+x,p.z+z)-anchor)&sr::kMask48,sr::kSlime};}
__global__ void rootsKernel(const sr::CellTerm* cells,uint32_t positions,int count,Job* out,unsigned long long* total){
    uint64_t t=uint64_t(blockIdx.x)*blockDim.x+threadIdx.x;if(t>=(uint64_t(positions)<<18))return;
    uint32_t pos=t>>18,root=t&262143;const auto* p=cells+uint64_t(pos)*count;
    for(int i=0;i<count;++i)if((((root+uint32_t(p[i].term))^uint32_t(sr::kCombinedXor))*uint32_t(sr::kJavaMultiplier)+11U)&(1U<<17))return;
    out[atomicAdd(total,1ULL)]={pos,root};
}
__global__ void prepareKernel(const sr::CellTerm* cells,const Job* jobs,uint32_t n,int count,uint32_t* prepared){
    uint32_t k=blockIdx.x*blockDim.x+threadIdx.x;if(k>=n)return;const auto job=jobs[k];
    for(int i=0;i<count;++i){uint64_t a=job.root+cells[uint64_t(job.position)*count+i].term;
        prepared[(uint64_t(k)*count+i)*2]=a>>18;
        prepared[(uint64_t(k)*count+i)*2+1]=((((a^sr::kCombinedXor)&262143)*sr::kJavaMultiplier+11)>>17);
    }
}
__host__ __device__ __forceinline__ uint32_t q(const uint32_t*p,int i,uint32_t m){return (p[2*i+1]+2U*uint32_t(sr::kJavaMultiplier)*((m+p[2*i])^uint32_t(sr::kCombinedXor>>18)))&0x7fffffffU;}
__host__ __device__ uint32_t groupMask(const uint32_t*p,uint32_t base,int count){
    uint32_t mask=65535,q0=q(p,0,base),qe=q(p,0,base+240);
    for(int i=1;i<count&&mask;++i){uint32_t cut=(256-((base+p[2*i])&255)+15)>>4,lo=(1U<<cut)-1;
        uint32_t r0=((q(p,i,base)-q0)&0x7fffffffU)%10,r1=((q(p,i,base+240)-qe)&0x7fffffffU)%10;
        mask&=((r0==0||r0==8)?lo:0U)|((r1==0||r1==8)?(65535U^lo):0U);
    }
    return mask;
}
struct Emit{
    Coord p;Match* out;uint64_t capacity;unsigned long long* total;
    __device__ void operator()(uint64_t seed){uint64_t slot=atomicAdd(total,1ULL);if(slot<capacity)out[slot]={(seed-sr::coordinateTerm(p.x,p.z))&sr::kMask48,p.x,p.z};}
};
__global__ void searchKernel(const sr::CellTerm*cells,const Coord*coords,const Job*jobs,uint32_t n,int count,const uint32_t*prepared,Match*out,uint64_t capacity,unsigned long long*total){
    uint64_t index=uint64_t(blockIdx.x)*blockDim.x+threadIdx.x;if(index>=(uint64_t(n)<<13))return;
    uint32_t g=index&8191,k=index>>13,base=((g>>4)<<8)+(g&15);const uint32_t* p=prepared+uint64_t(k)*count*2;
    uint32_t mask=groupMask(p,base,count);
    Job job=jobs[k];Emit emit{coords[job.position],out,capacity,total};
    while(mask){uint32_t j=__ffs(mask)-1;mask&=mask-1;fs::solveIntervals(job.root+(uint64_t(base+16*j)<<18),cells+uint64_t(job.position)*count,count,emit);}
}
template<class T>struct Buffer{T* p=nullptr;explicit Buffer(size_t n){if(n)CUDA(cudaMalloc(&p,n*sizeof(T)));}~Buffer(){if(p)cudaFree(p);}Buffer(const Buffer&)=delete;};
struct Metrics{int device;std::string name;uint64_t positions=0,classes=0,cacheHits=0,jobs=0,matches=0;double setupMs=0,scanMs=0,rootMs=0,searchMs=0;Clock::time_point start,end;std::vector<Match> sample;};
struct Shared{std::atomic<uint64_t> next{0},done{0};std::atomic<int> finished{0};int ready=0,workerCount=0;std::condition_variable cv;std::atomic<bool> cancel{false};std::mutex mutex;std::string error;};
void worker(const Config& c,int device,Shared& shared,Metrics& result){try{
    auto setup=Clock::now();CUDA(cudaSetDevice(device));cudaDeviceProp prop{};CUDA(cudaGetDeviceProperties(&prop,device));result.device=device;result.name=prop.name;
    int count=c.size*c.size;size_t rootsCapacity=size_t(c.positionBatch)*(1U<<17);
    Buffer<sr::CellTerm> dcells(size_t(c.positionBatch)*count);Buffer<Coord>dcoords(c.positionBatch);Buffer<Job>droots(rootsCapacity),djobs(c.jobBatch);
    Buffer<uint32_t>prep(size_t(c.jobBatch)*count*2);Buffer<unsigned long long>total(1);Buffer<Match>matches(c.maxResults);
    std::map<uint64_t,std::vector<uint32_t>> cache;std::deque<uint64_t> fifo;size_t cachedRoots=0;
    result.setupMs=std::chrono::duration<double,std::milli>(Clock::now()-setup).count();{
        std::unique_lock<std::mutex> lock(shared.mutex);++shared.ready;shared.cv.notify_all();
        shared.cv.wait(lock,[&](){return shared.ready==shared.workerCount||shared.cancel.load();});
        if(shared.cancel)return;
    }
    result.start=Clock::now();
    const uint64_t end=anchors(c.radius);
    while(!shared.cancel){uint64_t first=shared.next.fetch_add(c.positionBatch);if(first>=end)break;
        uint32_t num=std::min<uint64_t>(c.positionBatch,end-first);std::vector<Coord> coords;std::vector<uint64_t> keys;
        std::map<uint64_t,std::vector<uint32_t>> lists;std::vector<Coord> missing;std::vector<uint64_t> missingKeys;
        for(uint32_t i=0;i<num;++i){Coord p=coordinate(c,first+i);coords.push_back(p);uint64_t key=classKey(p);keys.push_back(key);
            if(lists.count(key)){++result.cacheHits;continue;}auto it=cache.find(key);
            if(it!=cache.end()){lists[key]=it->second;++result.cacheHits;}else{lists[key]={};missing.push_back(p);missingKeys.push_back(key);}
        }
        if(!missing.empty()){
            auto start=Clock::now();std::vector<sr::CellTerm> cells(missing.size()*count);for(size_t i=0;i<missing.size();++i)terms(missing[i],c.size,cells.data()+i*count);
            CUDA(cudaMemcpy(dcells.p,cells.data(),cells.size()*sizeof(sr::CellTerm),cudaMemcpyHostToDevice));CUDA(cudaMemset(total.p,0,sizeof(*total.p)));
            rootsKernel<<<unsigned((uint64_t(missing.size())<<18)/256),256>>>(dcells.p,missing.size(),count,droots.p,total.p);CUDA(cudaGetLastError());
            unsigned long long found;CUDA(cudaMemcpy(&found,total.p,sizeof(found),cudaMemcpyDeviceToHost));if(found>rootsCapacity)throw std::runtime_error("root capacity invariant failed");
            std::vector<Job> roots(found);if(found)CUDA(cudaMemcpy(roots.data(),droots.p,found*sizeof(Job),cudaMemcpyDeviceToHost));
            for(auto job:roots)lists[missingKeys[job.position]].push_back(job.root);
            result.classes+=missing.size();
            for(auto key:missingKeys){auto& list=lists[key];std::sort(list.begin(),list.end());if(c.cacheClasses){
                while(!fifo.empty()&&(cache.size()>=c.cacheClasses||cachedRoots+list.size()>(1U<<20))){auto old=fifo.front();fifo.pop_front();cachedRoots-=cache.at(old).size();cache.erase(old);}
                cache[key]=list;cachedRoots+=list.size();fifo.push_back(key);
            }}
            result.rootMs+=std::chrono::duration<double,std::milli>(Clock::now()-start).count();
        }
        // Construct concrete positions only for classes with roots.
        std::vector<Coord> active;std::vector<sr::CellTerm> cells;std::vector<Job> jobs;jobs.reserve(c.jobBatch);
        for(uint32_t i=0;i<num;++i)if(!lists.at(keys[i]).empty())active.push_back(coords[i]);
        if(!active.empty()){
            cells.resize(active.size()*count);for(size_t i=0;i<active.size();++i)terms(active[i],c.size,cells.data()+i*count);
            CUDA(cudaMemcpy(dcells.p,cells.data(),cells.size()*sizeof(sr::CellTerm),cudaMemcpyHostToDevice));CUDA(cudaMemcpy(dcoords.p,active.data(),active.size()*sizeof(Coord),cudaMemcpyHostToDevice));
            auto search=[&](){if(jobs.empty()||shared.cancel)return;auto start=Clock::now();
                CUDA(cudaMemcpy(djobs.p,jobs.data(),jobs.size()*sizeof(Job),cudaMemcpyHostToDevice));CUDA(cudaMemset(total.p,0,sizeof(*total.p)));
                prepareKernel<<<unsigned((jobs.size()+127)/128),128>>>(dcells.p,djobs.p,jobs.size(),count,prep.p);
                searchKernel<<<unsigned((uint64_t(jobs.size())<<13)/256),256>>>(dcells.p,dcoords.p,djobs.p,jobs.size(),count,prep.p,matches.p,c.maxResults,total.p);CUDA(cudaGetLastError());
                unsigned long long found;CUDA(cudaMemcpy(&found,total.p,sizeof(found),cudaMemcpyDeviceToHost));if(found>UINT64_MAX-result.matches)throw std::runtime_error("match count exceeds uint64");result.matches+=found;result.jobs+=jobs.size();
                size_t take=std::min<uint64_t>(found,c.maxResults-result.sample.size());if(take){size_t old=result.sample.size();result.sample.resize(old+take);CUDA(cudaMemcpy(result.sample.data()+old,matches.p,take*sizeof(Match),cudaMemcpyDeviceToHost));}
                result.searchMs+=std::chrono::duration<double,std::milli>(Clock::now()-start).count();jobs.clear();};
            for(uint32_t i=0;i<active.size()&&!shared.cancel;++i)for(uint32_t root:lists.at(classKey(active[i]))){jobs.push_back({i,root});if(jobs.size()==size_t(c.jobBatch))search();}
            search();
        }
        result.positions+=num;shared.done.fetch_add(num);
    }
    result.end=Clock::now();result.scanMs=std::chrono::duration<double,std::milli>(result.end-result.start).count();
}catch(const std::exception& e){shared.cancel=true;std::lock_guard<std::mutex> lock(shared.mutex);shared.error="GPU "+std::to_string(device)+": "+e.what();shared.cv.notify_all();}}
std::string escape(const std::string&s){std::ostringstream o;for(unsigned char ch:s){if(ch=='"'||ch=='\\')o<<'\\'<<char(ch);else if(ch<32)o<<"\\u"<<std::hex<<std::setw(4)<<std::setfill('0')<<int(ch)<<std::dec;else o<<char(ch);}return o.str();}
void help(){std::cout<<R"(slime-bench --radius R [options]
Search all effective 48-bit seeds for solid n x n slime chunk squares.
Ordinary candidate branch only; rejection-tail supplement is NOT included.

  --x X --z Z          Center chunk coordinates (default 0,0)
  --radius R           Manhattan radius in chunks (required, >=0)
  --size N             Square side, 1..6 (default 5)
  --devices 0,1,...     CUDA device indices, or all (default 0)
  --repeat N           Complete passes; fresh root cache each pass (default 1)
  --position-batch N   Anchors per scheduled batch, 1..256 (default 32)
  --job-batch N        Low18-root jobs per GPU launch, 1..65536 (default 2048)
  --cache-classes N    Per-GPU bounded FIFO root cache entries (default 4096)
  --max-results N      Save up to N example matches; count ALL (default 0)
  --output FILE        Also save the complete JSON report (file must not exist)
  --dry-run            Print scope JSON without using CUDA
  --list-devices       List visible CUDA devices
  --help               Show this help

Scope: abs(anchor_x-X)+abs(anchor_z-Z)<=R, including boundary.
Anchor is the minimum-x/minimum-z corner; the square may extend outside
this diamond. Coordinates are not clipped to Minecraft's world border.
Small squares can have enormous numbers of matching seeds and take very long.
)";}
int main(int argc,char**argv){try{
    Config c;bool radiusSet=false;for(int i=1;i<argc;++i){std::string a=argv[i];if(a=="--help"||a=="-h"){help();return 0;}if(a=="--dry-run"){c.dry=true;continue;}if(a=="--list-devices"){c.list=true;continue;}
        if(i+1>=argc)throw std::runtime_error("missing value for "+a);std::string v=argv[++i];
        auto bounded=[&](uint64_t lo,uint64_t hi){auto n=number(v);if(n<lo||n>hi)throw std::runtime_error("out of range: "+a);return n;};
        if(a=="--x")c.x=signed32(v);else if(a=="--z")c.z=signed32(v);else if(a=="--radius"){c.radius=bounded(0,INT32_MAX);radiusSet=true;}
        else if(a=="--size")c.size=bounded(1,6);else if(a=="--devices")c.devices=v;else if(a=="--repeat")c.repeats=bounded(1,1000);
        else if(a=="--position-batch")c.positionBatch=bounded(1,256);else if(a=="--job-batch")c.jobBatch=bounded(1,65536);
        else if(a=="--cache-classes")c.cacheClasses=bounded(0,1000000);else if(a=="--max-results")c.maxResults=bounded(0,1000000);
        else if(a=="--output")c.output=v;else throw std::runtime_error("unknown option: "+a);
    }
    if(c.list){int n;CUDA(cudaGetDeviceCount(&n));for(int i=0;i<n;++i){cudaDeviceProp p{};CUDA(cudaGetDeviceProperties(&p,i));std::cout<<i<<"\t"<<p.name<<"\tsm_"<<p.major<<p.minor<<"\t"<<p.totalGlobalMem/(1ULL<<20)<<" MiB\n";}return 0;}
    if(!radiusSet)throw std::runtime_error("--radius is required");
    for(int64_t center:{int64_t(c.x),int64_t(c.z)})if(center-int64_t(c.radius)<INT32_MIN||center+int64_t(c.radius)+c.size-1>INT32_MAX)throw std::runtime_error("coordinates including square must fit int32");
    if(!c.output.empty()&&std::ifstream(c.output).good())throw std::runtime_error("output already exists: "+c.output);
    std::ostringstream report;report<<std::setprecision(12)<<"{\"schema_version\":1,\"mode\":\"ordinary_all_effective_seeds\",\"rejection_tail_included\":false,\"center\":["<<c.x<<","<<c.z<<"],\"radius_chunks\":"<<c.radius<<",\"square_size\":"<<c.size<<",\"anchors\":"<<anchors(c.radius)<<",\"effective_seed_bits\":48,\"dry_run\":"<<(c.dry?"true":"false");
    if(c.dry){auto first=coordinate(c,0),last=coordinate(c,anchors(c.radius)-1);report<<",\"first_anchor\":["<<first.x<<","<<first.z<<"],\"last_anchor\":["<<last.x<<","<<last.z<<"]}";}
    else{
        int available;CUDA(cudaGetDeviceCount(&available));std::vector<int> devices;if(c.devices=="all"){for(int i=0;i<available;++i)devices.push_back(i);}else{
            std::stringstream ss(c.devices);std::string v;while(std::getline(ss,v,',')){auto id=number(v);if(id>=uint64_t(available))throw std::runtime_error("CUDA device unavailable: "+v);devices.push_back(id);}
        }
        if(devices.empty()||std::set<int>(devices.begin(),devices.end()).size()!=devices.size())throw std::runtime_error("device list must be nonempty and unique");
        report<<",\"position_batch\":"<<c.positionBatch<<",\"job_batch\":"<<c.jobBatch<<",\"cache_classes\":"<<c.cacheClasses<<",\"cuda_runtime_version\":"<<CUDART_VERSION<<",\"rounds\":[";
        std::vector<double> timings;uint64_t expected=0;
        for(int round=0;round<c.repeats;++round){Shared shared;shared.workerCount=devices.size();std::vector<Metrics> metrics(devices.size());std::vector<std::thread> workers;auto start=Clock::now();
            for(size_t i=0;i<devices.size();++i)workers.emplace_back([&,i](){worker(c,devices[i],shared,metrics[i]);++shared.finished;shared.cv.notify_all();});
            {std::unique_lock<std::mutex> lock(shared.mutex);while(shared.finished<int(devices.size())){
                if(!shared.cv.wait_for(lock,std::chrono::seconds(5),[&](){return shared.finished==int(devices.size());}))
                    std::cerr<<"progress: "<<shared.done<<"/"<<anchors(c.radius)<<" anchors (round "<<round+1<<")\n";
            }}
            for(auto&t:workers)t.join();
            if(!shared.error.empty())throw std::runtime_error(shared.error);
            double elapsed=std::chrono::duration<double,std::milli>(Clock::now()-start).count();auto first=metrics[0].start,last=metrics[0].end;
            uint64_t visited=0,jobs=0,matches=0;std::vector<Match> sample;
            for(auto&m:metrics){first=std::min(first,m.start);last=std::max(last,m.end);visited+=m.positions;jobs+=m.jobs;if(m.matches>UINT64_MAX-matches)throw std::runtime_error("total match count exceeds uint64");matches+=m.matches;for(auto hit:m.sample)if(sample.size()<c.maxResults)sample.push_back(hit);}
            if(visited!=anchors(c.radius))throw std::runtime_error("incomplete anchor coverage");if(round&&matches!=expected)throw std::runtime_error("repeat match count mismatch");expected=matches;
            double scan=std::chrono::duration<double,std::milli>(last-first).count();timings.push_back(scan);
            if(round)report<<",";report<<"{\"round\":"<<round+1<<",\"elapsed_ms\":"<<elapsed<<",\"scan_wall_ms\":"<<scan<<",\"anchors_per_second\":"<<visited/(scan/1000)<<",\"root_jobs\":"<<jobs<<",\"matches\":"<<matches<<",\"sample_truncated\":"<<(matches>sample.size()?"true":"false")<<",\"sample\":[";
            for(size_t i=0;i<sample.size();++i){if(i)report<<",";auto m=sample[i];report<<"["<<m.seed<<","<<m.x<<","<<m.z<<"]";}report<<"],\"gpus\":[";
            for(size_t i=0;i<metrics.size();++i){if(i)report<<",";auto&m=metrics[i];report<<"{\"device\":"<<m.device<<",\"name\":\""<<escape(m.name)<<"\",\"anchors\":"<<m.positions<<",\"root_classes_generated\":"<<m.classes<<",\"cache_hits\":"<<m.cacheHits<<",\"root_jobs\":"<<m.jobs<<",\"matches\":"<<m.matches<<",\"setup_ms\":"<<m.setupMs<<",\"scan_ms\":"<<m.scanMs<<",\"root_stage_ms\":"<<m.rootMs<<",\"search_stage_ms\":"<<m.searchMs<<"}";}report<<"]}";
            std::cerr<<"round "<<round+1<<"/"<<c.repeats<<": "<<visited<<" anchors, "<<matches<<" matches, scan "<<scan<<" ms, elapsed "<<elapsed<<" ms\n";
        }
        std::sort(timings.begin(),timings.end());double median=timings[timings.size()/2];if(timings.size()%2==0)median=(median+timings[timings.size()/2-1])/2;
        report<<"],\"median_scan_wall_ms\":"<<median<<"}";
    }
    if(!c.output.empty()){std::ofstream file(c.output);if(!file)throw std::runtime_error("cannot write report: "+c.output);file<<report.str()<<"\n";file.close();if(!file)throw std::runtime_error("failed writing report");}
    std::cout<<report.str()<<"\n";return 0;
}catch(const std::exception&e){std::cerr<<"error: "<<e.what()<<"\n";return 1;}}