"""
Command-Line Interface

CLI for OMOP cohort management via OHDSI WebAPI.

Author: Ankur Lohachab
Department of Advanced Computing Sciences, Maastricht University
"""

import argparse
import json
import sys
from pathlib import Path
from typing import Optional
from smart_omop.client import OMOPClient
from smart_omop.cohort import CohortBuilder, Gender
from smart_omop.heracles import HeraclesJobManager, HeraclesJobConfig
from smart_omop.fetcher import (
    fetch_cohort_data,
    create_and_generate_cohort,
    poll_generation_status
)


def cmd_create_cohort(args):
    """Create cohort definition."""
    # Build cohort
    builder = CohortBuilder(args.name, args.description or "")

    if args.concept_ids:
        concept_ids = [int(cid) for cid in args.concept_ids.split(',')]
        builder.with_condition("Primary Condition", concept_ids)

    if args.age_gte is not None or args.age_lte is not None:
        builder.with_age_range(args.age_gte, args.age_lte)

    if args.gender:
        gender_map = {
            'male': Gender.MALE,
            'female': Gender.FEMALE
        }
        builder.with_gender(gender_map[args.gender.lower()])

    cohort_def = builder.build()

    # Create via WebAPI
    with OMOPClient(args.base_url) as client:
        created = client.create_cohort(cohort_def.to_dict())

        print(f"✓ Cohort created successfully")
        print(f"  ID: {created['id']}")
        print(f"  Name: {created['name']}")

        if args.output:
            with open(args.output, 'w') as f:
                json.dump(created, f, indent=2)
            print(f"  Saved to: {args.output}")

        return created['id']


def cmd_generate(args):
    """Generate cohort on data source."""
    with OMOPClient(args.base_url) as client:
        result = client.generate_cohort(args.cohort_id, args.source_key)

        print(f"✓ Cohort generation started")
        print(f"  Cohort ID: {args.cohort_id}")
        print(f"  Source: {args.source_key}")

        if args.output:
            with open(args.output, 'w') as f:
                json.dump(result, f, indent=2)


def cmd_heracles(args):
    """Run Heracles characterization."""
    with OMOPClient(args.base_url) as client:
        # Create job manager
        manager = HeraclesJobManager(client)

        # Determine analysis IDs
        analysis_ids = None
        if args.analysis_ids:
            analysis_ids = [int(aid) for aid in args.analysis_ids.split(',')]

        # Create job configuration
        job_config = manager.create_job(
            cohort_ids=[args.cohort_id],
            source_key=args.source_key,
            job_name=args.job_name,
            analysis_ids=analysis_ids,
            small_cell_count=args.small_cell_count or 0,
            run_heel=args.run_heel,
            cohort_period_only=args.cohort_period_only
        )

        print(f"✓ Submitting Heracles characterization job")
        print(f"  Cohort ID: {args.cohort_id}")
        print(f"  Source: {args.source_key}")
        print(f"  Analysis IDs: {len(job_config.analysis_ids)} analyses")

        # Submit job
        result = manager.submit_job(
            job_config,
            poll=args.poll,
            timeout=args.timeout or 1800
        )

        print(f"✓ Job submitted successfully")

        exec_id = result.get('executionId') or result.get('id') or result.get('jobId')
        if exec_id:
            print(f"  Execution ID: {exec_id}")

        if args.poll:
            status = result.get('status') or result.get('exitStatus')
            print(f"  Final Status: {status}")


def cmd_results(args):
    """Fetch cohort results."""
    data = fetch_cohort_data(
        args.base_url,
        args.cohort_id,
        args.source_key,
        include_results=True
    )

    print(f"✓ Cohort data retrieved")
    print(f"  Cohort: {data['definition'].get('name', 'Unknown')}")

    if data.get('results'):
        print(f"  Results available: Yes")

        # Output format
        if args.output_format == 'json':
            output_data = data
        elif args.output_format == 'md':
            # Convert to markdown
            output_data = format_results_markdown(data)
        else:
            output_data = data

        if args.output:
            if args.output_format == 'md':
                with open(args.output, 'w') as f:
                    f.write(output_data)
            else:
                with open(args.output, 'w') as f:
                    json.dump(output_data, f, indent=2)
            print(f"  Saved to: {args.output}")
        else:
            if args.output_format == 'json':
                print(json.dumps(output_data, indent=2))
            else:
                print(output_data)
    else:
        print(f"  Results available: No")
        if data.get('results_error'):
            print(f"  Error: {data['results_error']}")


def cmd_get_report(args):
    """Get specific Heracles report."""
    with OMOPClient(args.base_url) as client:
        try:
            print(f"Fetching {args.type} report...")
            report_data = client.get_heracles_report(
                args.cohort_id,
                args.source_key,
                args.type,
                refresh=args.refresh
            )

            if args.output:
                with open(args.output, 'w') as f:
                    json.dump(report_data, f, indent=2)
                print(f"✓ Saved to {args.output}")
            else:
                print(json.dumps(report_data, indent=2))

        except Exception as e:
            print(f"✗ Failed: {e}", file=sys.stderr)
            sys.exit(1)


def cmd_export_reports(args):
    """Export all Heracles reports for a cohort."""
    with OMOPClient(args.base_url) as client:
        # Define report types to export
        report_types = [
            'dashboard',
            'person',
            'condition',
            'drug',
            'procedure',
            'measurement',
            'observation',
            'death'
        ]

        # Add components_summary if requested
        if args.include_components:
            report_types.append('components_summary')

        # Create output directory if specified
        output_dir = Path(args.output_dir) if args.output_dir else Path('.')
        output_dir.mkdir(parents=True, exist_ok=True)

        print(f"✓ Exporting Heracles reports")
        print(f"  Cohort ID: {args.cohort_id}")
        print(f"  Source: {args.source_key}")
        print(f"  Output directory: {output_dir}")
        print(f"  Refresh: {args.refresh}")
        print()

        exported = []
        failed = []

        for report_type in report_types:
            try:
                print(f"  Fetching {report_type} report...", end=' ')
                report_data = client.get_heracles_report(
                    args.cohort_id,
                    args.source_key,
                    report_type,
                    refresh=args.refresh
                )

                # Create filename
                filename = f"cohort{args.cohort_id}_{args.source_key}_{report_type}.json"
                filepath = output_dir / filename

                # Save to file
                with open(filepath, 'w') as f:
                    json.dump(report_data, f, indent=2)

                print(f"✓ Saved to {filename}")
                exported.append(report_type)

            except Exception as e:
                print(f"✗ Failed: {e}")
                failed.append(report_type)

        print()
        print(f"✓ Export complete")
        print(f"  Successful: {len(exported)}/{len(report_types)}")
        if exported:
            print(f"  Exported: {', '.join(exported)}")
        if failed:
            print(f"  Failed: {', '.join(failed)}")


def format_results_markdown(data: dict) -> str:
    """Format results as markdown."""
    lines = [
        f"# Cohort Results: {data['definition'].get('name', 'Unknown')}",
        "",
        f"**Cohort ID**: {data['cohort_id']}",
        f"**Source**: {data['source_key']}",
        ""
    ]

    if data.get('results'):
        results = data['results']
        lines.extend([
            "## Summary",
            "",
            f"- **Person Count**: {results.get('person_count', 'N/A')}",
            ""
        ])

    return "\n".join(lines)


def main():
    """Main CLI entry point."""
    parser = argparse.ArgumentParser(
        prog='smart-omop',
        description='OMOP CDM cohort management via OHDSI WebAPI'
    )

    parser.add_argument(
        '--base-url',
        required=True,
        help='WebAPI base URL (e.g., http://host:8080/WebAPI)'
    )

    subparsers = parser.add_subparsers(dest='command', help='Command to execute')

    # create-cohort
    create_parser = subparsers.add_parser('create-cohort', help='Create cohort definition')
    create_parser.add_argument('--name', required=True, help='Cohort name')
    create_parser.add_argument('--concept-ids', help='Comma-separated concept IDs')
    create_parser.add_argument('--description', help='Cohort description')
    create_parser.add_argument('--age-gte', type=int, help='Minimum age')
    create_parser.add_argument('--age-lte', type=int, help='Maximum age')
    create_parser.add_argument('--gender', choices=['male', 'female'], help='Gender filter')
    create_parser.add_argument('--output', help='Output file path')

    # generate
    gen_parser = subparsers.add_parser('generate', help='Generate cohort')
    gen_parser.add_argument('--cohort-id', type=int, required=True, help='Cohort ID')
    gen_parser.add_argument('--source-key', required=True, help='Source database key')
    gen_parser.add_argument('--output', help='Output file path')

    # heracles
    her_parser = subparsers.add_parser('heracles', help='Run Heracles characterization')
    her_parser.add_argument('--cohort-id', type=int, required=True, help='Cohort ID')
    her_parser.add_argument('--source-key', required=True, help='Source database key')
    her_parser.add_argument('--job-name', help='Optional job name for tracking')
    her_parser.add_argument('--analysis-ids', help='Comma-separated analysis IDs (default: comprehensive set)')
    her_parser.add_argument('--small-cell-count', type=int, help='Suppress cells below this count')
    her_parser.add_argument('--run-heel', action='store_true', help='Run Heracles Heel data quality checks')
    her_parser.add_argument('--cohort-period-only', action='store_true', help='Limit analyses to cohort period only')
    her_parser.add_argument('--poll', action='store_true', help='Poll until complete')
    her_parser.add_argument('--timeout', type=int, help='Poll timeout in seconds')

    # results
    res_parser = subparsers.add_parser('results', help='Fetch cohort results')
    res_parser.add_argument('--cohort-id', type=int, required=True, help='Cohort ID')
    res_parser.add_argument('--source-key', required=True, help='Source database key')
    res_parser.add_argument('--output', help='Output file path')
    res_parser.add_argument('--output-format', choices=['json', 'md'], default='json', help='Output format')

    # export-reports
    export_parser = subparsers.add_parser('export-reports', help='Export all Heracles reports')
    export_parser.add_argument('--cohort-id', type=int, required=True, help='Cohort ID')
    export_parser.add_argument('--source-key', required=True, help='Source database key')
    export_parser.add_argument('--output-dir', default='.', help='Output directory for report files')
    export_parser.add_argument('--refresh', action='store_true', help='Refresh report data')
    export_parser.add_argument('--include-components', action='store_true', help='Include components_summary report')

    # get-report
    report_parser = subparsers.add_parser('get-report', help='Get specific Heracles report')
    report_parser.add_argument('--cohort-id', type=int, required=True, help='Cohort ID')
    report_parser.add_argument('--source-key', required=True, help='Source database key')
    report_parser.add_argument('--type', required=True,
                               choices=['dashboard', 'person', 'condition', 'drug', 'procedure',
                                       'measurement', 'observation', 'death', 'components_summary'],
                               help='Report type')
    report_parser.add_argument('--output', help='Output file (default: stdout)')
    report_parser.add_argument('--refresh', action='store_true', help='Refresh report data')

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(1)

    # Execute command
    try:
        if args.command == 'create-cohort':
            cmd_create_cohort(args)
        elif args.command == 'generate':
            cmd_generate(args)
        elif args.command == 'heracles':
            cmd_heracles(args)
        elif args.command == 'results':
            cmd_results(args)
        elif args.command == 'get-report':
            cmd_get_report(args)
        elif args.command == 'export-reports':
            cmd_export_reports(args)
    except Exception as e:
        print(f"✗ Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()
