import argparse
import ast
import asyncio
import base64
import json
import logging

logging.basicConfig(level=logging.INFO)
log = logging.getLogger('eyepop.example')

import os
import sys
from argparse import Namespace
from io import BytesIO
from typing import Any

from dotenv import load_dotenv
from PIL import Image
from pybars import Compiler
from pydantic import TypeAdapter
from webui import webui

from eyepop import EyePopSdk, Job
from eyepop.data.data_types import TranscodeMode
from eyepop.data.types.asset import Area, RectangleArea
from eyepop.worker.worker_types import (
    ComponentParams,
    ContourFinderComponent,
    ContourType,
    CropForward,
    ForwardComponent,
    FullForward,
    InferenceComponent,
    MotionDetectConfig,
    MotionModel,
    Pop,
    TrackingComponent,
)

load_dotenv()

script_dir = os.path.dirname(__file__)


pop_examples = {
    "vehicles": Pop(components=[
        InferenceComponent(
            ability='eyepop.vehicle:latest',
            categoryName="vehicles",
            confidenceThreshold=0.8,
            forward=CropForward(
                targets=[TrackingComponent(
                    maxAgeSeconds=5.0,
                    motionModel=MotionModel.CONSTANT_VELOCITY,
                    agnostic=True,
                    classHysteresis=True,
                ), InferenceComponent(
                    ability='eyepop.vehicle.licence-plate:latest',
                    topK=1,
                    forward=CropForward(
                        targets=[InferenceComponent(
                            ability='eyepop.text.recognize.landscape:latest',
                            categoryName="licence-plate"
                        )]
                    )
                )]
            )
        )
    ]),

    "person": Pop(components=[
        InferenceComponent(
            ability='eyepop.person:latest',
            categoryName="person"
        )
    ]),

    "2d-body-points": Pop(components=[
        InferenceComponent(
            ability='eyepop.person:latest',
            categoryName="person",
            forward=CropForward(
                maxItems=128,
                targets=[InferenceComponent(
                    ability='eyepop.person.2d-body-points:latest',
                    categoryName="2d-body-points",
                    confidenceThreshold=0.25
                )]
        ))
    ]),

    "faces": Pop(components=[
        InferenceComponent(
            ability='eyepop.person:latest',
            categoryName="person",
            forward=CropForward(
                maxItems=128,
                targets=[InferenceComponent(
                    ability='eyepop.person.face.short-range:latest',
                    categoryName="2d-face-points",
                    forward=CropForward(
                        boxPadding=1.5,
                        orientationTargetAngle=-90.0,
                        targets=[InferenceComponent(
                            ability='eyepop.person.face-mesh:latest',
                            categoryName="3d-face-mesh"
                        )]
                    )
                )]
            )
        )
    ]),

    "hands": Pop(components=[
        InferenceComponent(
            ability='eyepop.person:latest',
            categoryName="person",
            forward=CropForward(
                maxItems=128,
                boxPadding=0.25,
                targets=[InferenceComponent(
                    ability='eyepop.person.palm:latest',
                    forward=CropForward(
                        includeClasses=["hand circumference"],
                        orientationTargetAngle=-90.0,
                        targets=[InferenceComponent(
                            ability='eyepop.person.3d-hand-points:latest',
                            categoryName="3d-hand-points"
                        )]
                    )
                )]
            )
        )
    ]),

    "3d-body-points": Pop(components=[
        InferenceComponent(
            ability='eyepop.person:latest',
            categoryName="person",
            forward=CropForward(
                boxPadding=0.5,
                targets=[InferenceComponent(
                    ability='eyepop.person.pose:latest',
                    hidden=True,
                    forward=CropForward(
                        boxPadding=0.5,
                        orientationTargetAngle=-90.0,
                        targets=[InferenceComponent(
                            ability='eyepop.person.3d-body-points.heavy:latest',
                            categoryName="3d-body-points",
                            confidenceThreshold=0.25
                        )]
                    )
                )]
            )
        )
    ]),

    "text": Pop(components=[
        InferenceComponent(
            ability='eyepop.text:latest',
            categoryName="text",
            confidenceThreshold=0.7,
            forward=CropForward(
                maxItems=128,
                targets=[InferenceComponent(
                    ability='eyepop.text.recognize.landscape:latest',
                    confidenceThreshold=0.1
                )]
            )
        )
    ]),

    "sam1": Pop(components=[
        InferenceComponent(
            ability='eyepop.sam.small:latest',
            id=1,
            forward=FullForward(
                targets=[ContourFinderComponent(
                    contourType=ContourType.POLYGON,
                    areaThreshold=0.005
                )]
            )
        )
    ]),

    "sam2": Pop(components=[
        InferenceComponent(
            ability="eyepop.sam2.encoder.tiny:latest",
            id=1,
            hidden=True,
            forward=FullForward(
                targets=[InferenceComponent(
                    ability='eyepop.sam2.decoder:latest',
                    forward=FullForward(
                        targets=[ContourFinderComponent(
                            contourType=ContourType.POLYGON,
                            areaThreshold=0.005
                        )]
                    )
                )]
            )
        )
    ]),
    "image-contents": Pop(components=[
        InferenceComponent(
            id=1,
            ability='eyepop.image-contents:latest',
        )
    ]),
    "localize-objects": Pop(components=[
        InferenceComponent(
            id=1,
            ability='eyepop.localize-objects:latest',
        )
    ]),
    "localize-objects-plus": Pop(components=[
        InferenceComponent(
            id=1,
            ability='eyepop.localize-objects:latest',
            params={
                "prompts": [{"prompt": "person"}]
            },
            forward=CropForward(
                targets=[InferenceComponent(
                    ability='eyepop.image-contents:latest',
                    params={
                        "prompts": [{"prompt": "hair color blond"},{"prompt": "hair color brown"}]
                    }
                )],
            )
        )
    ]),
    "localize-objects-t4": Pop(components=[
        InferenceComponent(
            id=1,
            ability='eyepop.localize-objects:latest',
            params={
                "prompts": [{"prompt": "person"}]
            },
            forward=CropForward(
                targets=[InferenceComponent(
                    ability='eyepop.image-contents-t4:latest',
                    params={
                        "prompts": [{"prompt": "shirt color?"}]
                    }
                )],
            )
        )
    ]),
}

def list_of_points(arg: str) -> list[dict[str, Any]]:
    points = []
    points_as_tuples = ast.literal_eval(f'[{arg}]')
    for tuple in points_as_tuples:
        points.append({
            "x": tuple[0],
            "y": tuple[1]
        })
    return points


def list_of_boxes(arg: str) -> list[dict[str, Any]]:
    boxes = []
    boxes_as_tuples = ast.literal_eval(f'[{arg}]')
    for tuple in boxes_as_tuples:
        boxes.append({
            "topLeft": {
                "x": tuple[0],
                "y": tuple[1]
            },
            "bottomRight": {
                "x": tuple[2],
                "y": tuple[3]
            }
        })
    return boxes


def rectangle_roi(arg: str) -> Area:
    roi = ast.literal_eval(arg)

    return RectangleArea(
        x=roi[0],
        y=roi[1],
        width=roi[2],
        height=roi[3],
    )


def add_optional_tracking_to_component(component: ForwardComponent, tracking_args: Namespace):
    if tracking_args.tracking:
        tracking_component = TrackingComponent()
        if tracking_args.tracking_reid_model is not None:
            tracking_component.reidModel = tracking_args.tracking_reid_model
        if tracking_args.tracking_max_age is not None:
            tracking_component.maxAgeSeconds = tracking_args.tracking_max_age
        if tracking_args.tracking_iou_threshold is not None:
            tracking_component.iouThreshold = tracking_args.tracking_iou_threshold
        if tracking_args.tracking_sim_threshold is not None:
            tracking_component.simThreshold = tracking_args.tracking_sim_threshold
        if tracking_args.tracking_agnostic is not None:
            tracking_component.agnostic = tracking_args.tracking_agnostic
        if tracking_args.tracking_motion_model is not None:
            tracking_component.motionModel = MotionModel(tracking_args.tracking_motion_model)
        component.forward = CropForward(
            targets=[tracking_component]
        )


parser = argparse.ArgumentParser(
                    prog='Pop examples',
                    description='Demonstrates the composition of a Pop',
                    epilog='.')
parser.add_argument('-d', '--dump', required=False, help="dump composable pop definition to stdout", default=False, action="store_true")
parser.add_argument('-l', '--local-path', required=False, type=str, default=False, help="run the inference on a local file, or all files on a directory")
parser.add_argument('-a', '--asset-uuid', required=False, type=str, default=False, help="run the inference on an asset by its Uuid")
parser.add_argument('-u', '--url', required=False, type=str, default=False, help="run the inference on a remote Url")
parser.add_argument('-p', '--pop', required=False, type=str, help="run this pop", choices=list(pop_examples.keys()))
parser.add_argument('-m', '--model-uuid', required=False, type=str, action="append", help="run this model(s) by uuid")
parser.add_argument('-ma', '--model-alias', required=False, type=str, action="append", help="run this model(s) by its tagged alias")
parser.add_argument('-ms1', '--model-uuid-sam1', required=False, type=str, help="run this model by uuid and compose with SAM1 (EfficientSAM) and Contour Finder")
parser.add_argument('-ms2', '--model-uuid-sam2', required=False, type=str, help="run this model by uuid and compose with SAM2 and Contour Finder")
parser.add_argument('-po', '--points', required=False, type=list_of_points, help="List of POIs as coordinates like (x1, y1), (x2, y2) in the original image coordinate system")
parser.add_argument('-bo', '--boxes', required=False, type=list_of_boxes, help="List of POIs as boxes like (left1, top1, right1, bottom1), (left1, top1, right1, bottom1) in the original image coordinate system")
parser.add_argument('-sp', '--single-prompt', required=False, type=str, help="Single prompt to pass as parameter")
parser.add_argument('-sl', '--single-label', required=False, type=str, help="Single label to use as result for single prompt to pass as parameter")
parser.add_argument('-pr', '--prompt', required=False, type=str, help="Prompt to pass as parameter", action="append")
parser.add_argument('-v', '--visualize', required=False, help="show rendered output", default=False, action="store_true")
parser.add_argument('-o', '--output', required=False, help="print results to stdout", default=False, action="store_true")
parser.add_argument('-ds', '--dataset-uuid', required=False, type=str, help="Ingest all assets into a dataset uuid", default=None)
parser.add_argument('-tk', '--top-k', required=False, type=int, help="For --model-uuid and -model-alias apply this top-k filter", default=None)
parser.add_argument('-ct', '--confidence-threshold', required=False, type=float, help="For --model-uuid and -model-alias apply this confidence threshold filter", default=None)
parser.add_argument('--fps', required=False, type=str, help="For --model-uuid and -model-alias apply this target fps - e.g. 15/1", default=None)
parser.add_argument('-vl', '--video-chunk-length-seconds', required=False, type=float, help="For --model-uuid and -model-alias apply this video chunk length", default=None)
parser.add_argument('-vo', '--video-chunk-overlap', required=False, type=float, help="For --model-uuid and -model-alias apply this video chunk overlap", default=None)

# Optional tracking for simple pops my model uuid oder model alias
parser.add_argument('--tracking', required=False, help="Track objects in videos", default=False, action="store_true")
parser.add_argument('--tracking-reid-model', required=False, help="Use re-id model uuid for tracking", default=None, type=str)
parser.add_argument('--tracking-agnostic', required=False, help="Track objects class-agnostic", default=False, action="store_true")
parser.add_argument('--tracking-max-age', required=False, help="Max age in seconds for unmatched tracks", default=None, type=float)
parser.add_argument('--tracking-iou-threshold', required=False, help="IoU threshold to match tracks", default=None, type=float)
parser.add_argument('--tracking-sim-threshold', required=False, help="Similarity threshold to match tracks by re-id", default=None, type=float)
parser.add_argument('--tracking-motion-model', required=False, help="Pick a motion model for tracking, one of 'random_walk', 'constant_velocity' or 'constant_acceleration'", default=None, type=str)

# Optional motion detection parameters
parser.add_argument('--motion-detect', required=False, help="Skip video frames w/o detected motion", default=False, action="store_true")

# Optional global ROI parameters
parser.add_argument('--roi', required=False, type=rectangle_roi, help="Rectangular ROI as (x, y, width, height)")


main_args = parser.parse_args()

if not main_args.local_path and not main_args.url and not main_args.asset_uuid:
    print("Need something to run inference on; pass either --url or --local-path or --asset-uuid")
    parser.print_help()
    sys.exit(1)

if not main_args.pop and not main_args.model_uuid and not main_args.model_alias and not main_args.model_uuid_sam1 and not main_args.model_uuid_sam2:
    print("Need something do do, pass either --pop or --model-uuid or --model-alias or --model-uuid-sam1 or --model-uuid-sam2")
    parser.print_help()
    sys.exit(1)

if main_args.pop:
    pop = pop_examples[main_args.pop]
elif main_args.model_uuid:
    pop = Pop(components=[
        InferenceComponent(
            id=i+1,
            abilityUuid=uuid,
            targetFps=main_args.fps,
        ) for i, uuid in enumerate(main_args.model_uuid)
    ])
    if main_args.top_k is not None:
        for c in pop.components:
            c.topK = main_args.top_k

    if main_args.confidence_threshold is not None:
        for c in pop.components:
            c.confidenceThreshold = main_args.confidence_threshold

    if main_args.video_chunk_length_seconds is not None:
        for c in pop.components:
            c.videoChunkLengthSeconds = main_args.video_chunk_length_seconds

    if main_args.video_chunk_overlap is not None:
        for c in pop.components:
            c.videoChunkOverlap = main_args.video_chunk_overlap

    add_optional_tracking_to_component(pop.components[0], main_args)
elif main_args.model_alias:
    pop = Pop(components=[
        InferenceComponent(
            id=i+1,
            ability=alias,
            targetFps = main_args.fps,
        ) for i, alias in enumerate(main_args.model_alias)
    ])

    if main_args.top_k is not None:
        for c in pop.components:
            c.topK = main_args.top_k

    if main_args.confidence_threshold is not None:
        for c in pop.components:
            c.confidenceThreshold = main_args.confidence_threshold

    if main_args.video_chunk_length_seconds is not None:
        for c in pop.components:
            c.videoChunkLengthSeconds = main_args.video_chunk_length_seconds

    if main_args.video_chunk_overlap is not None:
        for c in pop.components:
            c.videoChunkOverlap = main_args.video_chunk_overlap

    add_optional_tracking_to_component(pop.components[0], main_args)
elif main_args.model_uuid_sam1:
    pop = Pop(components=[
        InferenceComponent(
            abilityUuid=main_args.model_uuid_sam1,
            forward=CropForward(
                targets=[InferenceComponent(
                    ability='eyepop.sam.small:latest',
                    forward=FullForward(
                        targets=[ContourFinderComponent(
                            contourType=ContourType.POLYGON,
                            areaThreshold=0.005
                        )]
                    )
                )]
            )
        )
    ])
elif main_args.model_uuid_sam2:
    pop = Pop(components=[
        InferenceComponent(
            ability="eyepop.sam2.encoder.tiny:latest",
            hidden=True,
            forward=FullForward(
                targets=[InferenceComponent(
                    abilityUuid=main_args.model_uuid_sam2,
                    forward=CropForward(
                        targets=[InferenceComponent(
                            ability='eyepop.sam2.decoder:latest',
                            forward=FullForward(
                                targets=[ContourFinderComponent(
                                    contourType=ContourType.POLYGON,
                                    areaThreshold=0.005
                                )]
                            )
                        )]
                    )
                )]
            )
        )
    ])
else:
    raise ValueError("pop or model required")

params = None
if main_args.points:
    params = [
        ComponentParams(componentId=1, values={
          "roi": {
              "points": main_args.points
          }
        })
    ]
elif main_args.boxes:
    params = [
        ComponentParams(componentId=1, values={
            "roi": {
                "boxes": main_args.boxes
            }
        })
    ]
elif main_args.prompt is not None and len(main_args.prompt) > 0:
    params = [
        ComponentParams(componentId=1, values={
            "prompts": [{"prompt": p} for p in main_args.prompt]
        })
    ]
elif main_args.single_prompt is not None:
    params = [
        ComponentParams(componentId=1, values={
            "prompt": main_args.single_prompt,
            "label": main_args.single_label if main_args.single_label else main_args.single_prompt
        })
    ]

async def main(args) -> tuple[dict[str, Any] | None, str | None]:
    visualize_prediction = None
    visualize_path = None
    example_image_src = None
    motion_detect = MotionDetectConfig(motionGap=1) if args.motion_detect else None

    if args.dump:
        print("Pop:", pop.model_dump_json(indent=2))
        if params:
            print("Params:", TypeAdapter(list[ComponentParams]).dump_json(params, indent=2).decode("utf-8"))

    async with EyePopSdk.workerEndpoint(dataset_uuid=args.dataset_uuid, is_async=True) as endpoint:
        await endpoint.set_pop(pop)

        if args.local_path:
            if not os.path.exists(args.local_path):
                log.warning(f"local path {args.local_path} does not exist")
                sys.exit(1)
            if os.path.isfile(args.local_path):
                local_files = [args.local_path]
            else:
                local_files = []

                for f in os.listdir(args.local_path):
                    local_file = os.path.join(args.local_path, f)
                    if os.path.isfile(local_file):
                        local_files.append(local_file)
            jobs = []
            async def on_ready(job: Job, path: str):
                nonlocal visualize_prediction
                nonlocal visualize_path
                while result := await job.predict():
                    visualize_prediction = result
                    visualize_path = path
                    if args.output:
                        print(path, json.dumps(result, indent=2))
            for local_file in local_files:
                job = await endpoint.upload(local_file, params=params, motion_detect=motion_detect, roi=args.roi)
                jobs.append(on_ready(job, local_file))
            await asyncio.gather(*jobs)
            if args.visualize and visualize_prediction is not None:
                image = Image.open(visualize_path)
                buffer = BytesIO()
                image.save(buffer, format="PNG")
                example_image_src = f"data:image/png;base64, {base64.b64encode(buffer.getvalue()).decode()}"
        elif args.url:
            job = await endpoint.load_from(args.url, params=params, motion_detect=motion_detect, roi=args.roi)
            while result := await job.predict():
                visualize_prediction = result
                if args.output:
                    print(json.dumps(result, indent=2))
            if args.visualize:
                example_image_src = args.url
        elif args.asset_uuid:
            job = await endpoint.load_asset(args.asset_uuid, params=params, motion_detect=motion_detect, roi=args.roi)
            while result := await job.predict():
                visualize_prediction = result
                if args.output:
                    print(json.dumps(result, indent=2))
            if args.visualize:
                async with EyePopSdk.dataEndpoint(is_async=True) as dataEndpoint:
                    buffer = await dataEndpoint.download_asset(
                        args.asset_uuid,
                        transcode_mode=TranscodeMode.image_original_size
                    ).read()
                    example_image_src = f"data:image/jpeg;base64, {base64.b64encode(buffer).decode()}"
    return visualize_prediction, example_image_src

visualize_result, example_image_src = asyncio.run(main(main_args))
if main_args.visualize:
    with open(os.path.join(script_dir, 'viewer.html')) as file:
        compiler = Compiler()
        html_template = compiler.compile(file.read())

    preview = html_template({
        'image_src': example_image_src,
        'result_json': json.dumps(visualize_result)
    })
    window = webui.window()
    window.set_root_folder('.')
    window.show(preview)
    webui.wait()



