"""Portable corpus statistics, metadata groups and matched task analysis."""
from __future__ import annotations
import argparse
from collections import Counter, defaultdict
from itertools import product
import json
from pathlib import Path
from .common import read_jsonl, prediction_text
from .metrics import score_task1, score_task2, score_task3, score_task4, score_task5
from .parsers import parse_task5

def corpus_stats(rows):
    sources=set(); tiers=Counter(); density=Counter(); pairs=Counter(); rationale=Counter(); labels=defaultdict(Counter)
    for row in rows:
        sid=row['source_id']
        if sid in sources: raise ValueError(f'Duplicate source_id: {sid}')
        sources.add(sid)
        u=row['user_utterance']
        rationale.update(k for k,v in u.get('rationales',{}).items() if v and str(v).strip())
        for k,v in u.get('labels',{}).items():
            labels[k].update(set(v) if isinstance(v,list) else [str(v)])
        for tier,responses in row['responses'].items():
            if not isinstance(responses,list): continue
            tiers[tier]+=len(responses)
            if tier=='accepted':
                for response in responses:
                    strategies=set(response.get('strategies',[]))
                    density[len(strategies)]+=1
                    pairs.update(product(strategies,repeat=2))
    total=len(sources)
    return {'sources':total,'response_tiers':dict(tiers),'candidates':sum(tiers.values()),
        'rationale_coverage':{k:{'count':n,'rate':n/total if total else None} for k,n in sorted(rationale.items())},
        'accepted_strategy_density':{str(k):{'count':v,'rate':v/tiers['accepted'] if tiers['accepted'] else None} for k,v in sorted(density.items())},
        'accepted_strategy_cooccurrence':[{'left':a,'right':b,'count':n} for (a,b),n in sorted(pairs.items())],
        'label_counts':{k:dict(sorted(v.items())) for k,v in sorted(labels.items())}}

def dotted(row, key):
    value=row
    for part in key.split('.'):
        if not isinstance(value,dict): return None
        value=value.get(part)
    return value

def grouped_metrics(rows,task,fields):
    groups=defaultdict(list)
    for row in rows:
        choices=[]
        for field in fields:
            value=dotted(row,field)
            choices.append(sorted({str(x) for x in value}) or ['__missing__'] if isinstance(value,list) else [str(value) if value is not None else '__missing__'])
        for key in product(*choices): groups[key].append(row)
    scorer={1:score_task1,2:score_task2,3:score_task3,4:score_task4,5:score_task5}[task]
    return {'task':task,'group_by':fields,'groups':[{'values':dict(zip(fields,key)),'metrics':scorer(values)} for key,values in sorted(groups.items())]}

def matched_analysis(task3_paths,task5_paths):
    from .overlap import norm_model, source_id, table_for, summarize_group
    task3={}; task5=[]; seen=set(); t3_unclassified=0
    for path in task3_paths:
        for row in read_jsonl(path):
            model=norm_model(row.get('evaluated_model') or row.get('source_model'))
            sid=row.get('source_id') or source_id(row)
            if not model or not sid: raise ValueError('Task 3 requires source and model identifiers')
            measurement=row.get('measurement') or {}
            label=measurement.get('label')
            if 'classification_status' in row:
                label=('safe' if row['is_safe'] else 'unsafe') if row.get('classification_status')=='classified' and type(row.get('is_safe')) is bool else None
            if label not in ('safe','unsafe'):
                t3_unclassified+=1; continue
            key=(str(sid),model)
            if key in task3: raise ValueError(f'Duplicate Task 3 source/model key: {key}')
            task3[key]=label=='safe'
    for path in task5_paths:
        for row in read_jsonl(path):
            model=norm_model(row.get('evaluated_model') or row.get('model_name') or Path(path).stem)
            sid=row.get('source_id') or source_id(row)
            metadata=row.get('metadata') or {}
            variant=metadata.get('prompt_variant') or metadata.get('task7_prompt_variant')
            if not sid or variant not in ('safe_guided','tool_forced','unsafe_induced'): raise ValueError('Task 5 requires source ID and a recognized framing')
            key=(str(sid),model,variant)
            if key in seen: raise ValueError(f'Duplicate Task 5 source/model/framing key: {key}')
            seen.add(key)
            gold=row.get('gold') or {}
            if 'parsed' in row:
                parsed=row.get('parsed') or {}; success=row.get('status')=='success'
            else:
                parsed=parse_task5(prediction_text(row),gold.get('available_tools',[])) or {}; success=row.get('api_status')=='ok'
            decision=parsed.get('decision'); tool=parsed.get('tool_name')
            parsed_ok=bool(success and decision)
            unsafe=bool(parsed_ok and decision=='call_tool' and tool in set(gold.get('unsafe_tools') or []))
            task5.append({'source_id':str(sid),'model':model,'variant':variant,'parsed_ok':parsed_ok,'unsafe_call':unsafe,'decision':decision,'tool_name':tool})
    counters,matched=table_for(task5,task3)
    return {'task3_unique_source_model_pairs':len(task3),'task3_unclassified_rows':t3_unclassified,
        'task5_rows':len(task5),'matched_rows':len(matched),
        'matched_unique_source_model_pairs':len({(r['source_id'],r['model']) for r in matched}),
        'task5_rows_without_classified_task3_match':len(task5)-len(matched),
        'summary':{k:summarize_group(counters[k]) for k in ['safe_guided','tool_forced','unsafe_induced','all']}}

def main():
    p=argparse.ArgumentParser(description=__doc__); sub=p.add_subparsers(dest='command',required=True)
    stats=sub.add_parser('stats'); stats.add_argument('--corpus',required=True)
    groups=sub.add_parser('groups'); groups.add_argument('--task',type=int,choices=range(1,6),required=True); groups.add_argument('--input',nargs='+',required=True); groups.add_argument('--group-by',nargs='+',required=True)
    overlap=sub.add_parser('overlap'); overlap.add_argument('--task3',nargs='+',required=True); overlap.add_argument('--task5',nargs='+',required=True)
    report=sub.add_parser('report',help='Own-model diagnostic report (JSON, CSV, Markdown, HTML)')
    for task in range(1,6): report.add_argument(f'--task{task}',nargs='+')
    report.add_argument('--output-dir',required=True)
    report.add_argument('--min-sources',type=int,default=20)
    report.add_argument('--min-coverage',type=float,default=.9)
    report.add_argument('--max-cases',type=int,default=10)
    for command in [stats,groups,overlap]: command.add_argument('--output',required=True)
    a=p.parse_args()
    if a.command=='report':
        from .profile import write_report
        inputs={t:getattr(a,f'task{t}') for t in range(1,6) if getattr(a,f'task{t}')}
        if not inputs: p.error('Supply at least one --taskN prediction/measurement file')
        write_report(inputs,a.output_dir,a.min_sources,a.min_coverage,a.max_cases)
        print(Path(a.output_dir).resolve()); return
    if a.command=='stats': result=corpus_stats(read_jsonl(a.corpus))
    elif a.command=='groups': result=grouped_metrics([row for f in a.input for row in read_jsonl(f)],a.task,a.group_by)
    else: result=matched_analysis(a.task3,a.task5)
    out=Path(a.output); out.parent.mkdir(parents=True,exist_ok=True); out.write_text(json.dumps(result,ensure_ascii=False,indent=2)+'\n',encoding='utf-8')
    print(out)

if __name__=='__main__': main()