package producer

import (
	"encoding/json"
	"fmt"
	"image"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/airbusgeo/godal"
	"github.com/paulmach/orb"
	"github.com/paulmach/orb/geojson"
	log "github.com/sirupsen/logrus"
	"gocv.io/x/gocv"
	"starwiz.cn/sjy01/image-proc/pkg/config"
	"starwiz.cn/sjy01/image-proc/pkg/fusion"
	"starwiz.cn/sjy01/image-proc/pkg/payload"
	"starwiz.cn/sjy01/image-proc/pkg/rrc"
	"starwiz.cn/sjy01/image-proc/pkg/utils"
)

type Scene struct {
	Type       string
	Width      int
	Height     int
	X          int // coordinate of the left top corner of the scene
	Y          int
	Mat        []gocv.Mat
	Tiff       string
	Meta       *ProductMeta
	SceneId    string
	SceneIndex string
}

func (s *Scene) Cleanup() {
	for i := range s.Mat {
		s.Mat[i].Close()
	}
}

func CleanScenes(scenes []*Scene) {
	for _, scene := range scenes {
		scene.Cleanup()
	}
}

// 对 PAN 和 配准后的MSS 在 Y 方向进行分景，景之间有25%的重叠
// 默认分景大小：
// MSS  2336 * 2336 - 1764
// PAN  9344 * 9344 - 7056
func (r *Registrator) SubScenes() (panScenes []*Scene, mssScenes []*Scene, err error) {
	if len(r.Params.Targets.Targets) > 0 && r.Params.Targets.Targets[0].EndLine > 0 {
		return r.RoiScenes()
	}

	hPAN := (payload.PAN_PIXEL_WIDTH - payload.MSS_PIXEL_WIDTH)
	hPANOverlap := payload.MSS_PIXEL_WIDTH
	panScenesCnt := r.PanHeight / hPAN

	for i := 0; i <= panScenesCnt; i++ {
		y1 := (i+1)*hPAN + hPANOverlap
		if y1 > r.PanHeight {
			y1 = r.PanHeight
		}

		scene := &Scene{
			Type:       "PAN",
			Width:      payload.PAN_PIXEL_WIDTH,
			Height:     y1 - i*hPAN,
			X:          0,
			Y:          i * hPAN,
			SceneIndex: fmt.Sprintf("%03d", i+1),
		}

		if scene.Height < scene.Width/2 {
			log.Info("scene height too small, skip")
			continue
		}

		name := filepath.Base(r.Params.PanTiffFile)
		name = strings.TrimSuffix(name, ".tiff")
		scene.SceneId = fmt.Sprintf("%s_%03d", name, i+1)

		mat := r.PanImage.Region(image.Rect(0, i*hPAN, payload.PAN_PIXEL_WIDTH, y1))
		if config.GCONFIG.Radiation.PANRemoveHfNoise {
			log.Println("applying hf noise filter on", scene.SceneId)
			matFiltered := rrc.HFNoiseFilter(mat, float64(mat.Cols())*config.GCONFIG.Radiation.HfRadiusRatio)
			mat.Close()
			mat = matFiltered
			runtime.GC()
		}

		if config.GCONFIG.Radiation.SceneMomentMatching {
			rrc.DoMomentMatching(mat)
		}

		scene.Mat = append(scene.Mat, mat)
		panScenes = append(panScenes, scene)
	}

	hMSS := hPAN / 4
	hMSSOverlap := hPANOverlap / 4
	mssScenesCnt := r.MssHeight / hMSS
	for i := 0; i <= mssScenesCnt; i++ {
		y1 := (i+1)*hMSS + hMSSOverlap
		if y1 > r.MssHeight {
			y1 = r.MssHeight
		}

		scene := &Scene{
			Type:       "MSS",
			Width:      payload.MSS_PIXEL_WIDTH,
			Height:     y1 - i*hMSS,
			X:          0,
			Y:          i * hMSS,
			SceneIndex: fmt.Sprintf("%03d", i+1),
		}

		if scene.Height < scene.Width/2 {
			log.Info("scene height too small, skip")
			continue
		}

		name := filepath.Base(r.Params.MssTiffFile)
		name = strings.TrimSuffix(name, ".tiff")
		scene.SceneId = fmt.Sprintf("%s_%03d", name, i+1)

		for band := 0; band < 4; band++ {
			mat := r.registeredMssImages[band].Region(image.Rect(0, i*hMSS, 2336, y1))
			if config.GCONFIG.Radiation.MSSRemoveHfNoise {
				log.Println("applying hf noise filter on", scene.SceneId)
				matFiltered := rrc.HFNoiseFilter(mat, float64(mat.Cols())*config.GCONFIG.Radiation.HfRadiusRatio/4)
				mat.Close()
				mat = matFiltered
				runtime.GC()
			}

			if config.GCONFIG.Radiation.SceneMomentMatching {
				rrc.DoMomentMatching(mat)
			}

			scene.Mat = append(scene.Mat, mat)
		}

		mssScenes = append(mssScenes, scene)
	}

	if len(panScenes) != len(mssScenes) {
		log.Error("pan and mss scenes count not match")
		err = fmt.Errorf("pan and mss scenes count not match")
	}

	// var b bytes.Buffer
	// gocv.MatProfile.WriteTo(&b, 1)
	// fmt.Print(b.String())

	return panScenes, mssScenes, err
}

func (r *Registrator) OutputL1A(panScenes []*Scene, mssScenes []*Scene) error {
	var fc geojson.FeatureCollection
	for i, scene := range panScenes {
		dir := filepath.Join(r.Params.OutputDir, scene.SceneIndex, "PAN")
		os.MkdirAll(dir, 0755)

		filename := fmt.Sprintf("%s_L1A.tiff", scene.SceneId)
		scene.Tiff = filepath.Join(dir, filename)
		scene.Meta = r.makeProductMeta(scene)
		r.ComputeMetaAndRPC(scene)

		err := utils.SavePanToGDALGTiff(scene.Mat[0],
			scene.Meta.Corners.UpperLeft.Longitude,
			scene.Meta.Corners.UpperLeft.Latitude,
			scene.Tiff,
			scene.Meta.Gsd)
		if err != nil {
			log.Errorf("save pan scene %d to tiff failed: %v", i+1, err)
			return err
		}

		scene.Meta.DataSize = sizeOfFile(scene.Tiff)
		metaFile := strings.Replace(scene.Tiff, ".tiff", ".meta.xml", 1)
		writeProductMeta(scene.Meta, metaFile)
		jpg := strings.Replace(scene.Tiff, ".tiff", ".jpg", 1)
		GTiffToJPG(scene.Tiff, jpg, "PAN", false)
		r.report.Scenes = append(r.report.Scenes, ReportScene{
			Name:        scene.SceneId,
			TiffData:    scene.Tiff,
			MetaData:    metaFile,
			BrowserData: jpg,
		})

		feature := geojson.NewFeature(orb.Polygon{
			{
				{scene.Meta.Corners.UpperLeft.Longitude, scene.Meta.Corners.UpperLeft.Latitude},
				{scene.Meta.Corners.UpperRight.Longitude, scene.Meta.Corners.UpperRight.Latitude},
				{scene.Meta.Corners.LowerRight.Longitude, scene.Meta.Corners.LowerRight.Latitude},
				{scene.Meta.Corners.LowerLeft.Longitude, scene.Meta.Corners.LowerLeft.Latitude},
				{scene.Meta.Corners.UpperLeft.Longitude, scene.Meta.Corners.UpperLeft.Latitude},
			},
		})
		fc.Features = append(fc.Features, feature)
	}

	data, _ := json.Marshal(fc)
	// 输出 GeoJSON 数据
	log.Debug(string(data))

	for _, scene := range mssScenes {
		dir := filepath.Join(r.Params.OutputDir, scene.SceneIndex, "MSS")
		os.MkdirAll(dir, 0755)

		filename := fmt.Sprintf("%s_L1A.tiff", scene.SceneId)
		scene.Tiff = filepath.Join(dir, filename)
		scene.Meta = r.makeProductMeta(scene)
		r.ComputeMetaAndRPC(scene)

		rgbirImage, _ := r.MergeMSSToBGRNIR(scene.Mat)
		defer rgbirImage.Close()

		err := utils.SaveBGRToGDALGTiff(rgbirImage,
			4,
			scene.Meta.Corners.UpperLeft.Longitude,
			scene.Meta.Corners.UpperLeft.Latitude,
			scene.Meta.Gsd,
			[]godal.ColorInterp{godal.CIBlue, godal.CIGreen, godal.CIRed, godal.CIUndefined},
			scene.Tiff)
		if err != nil {
			return err
		}

		scene.Meta.DataSize = sizeOfFile(scene.Tiff)
		metaFile := strings.Replace(scene.Tiff, ".tiff", ".meta.xml", 1)
		writeProductMeta(scene.Meta, metaFile)
		jpg := strings.Replace(scene.Tiff, ".tiff", ".jpg", 1)
		GTiffToJPG(scene.Tiff, jpg, "MSS", false)
		r.report.Scenes = append(r.report.Scenes, ReportScene{
			Name:        scene.SceneId,
			TiffData:    scene.Tiff,
			MetaData:    metaFile,
			BrowserData: jpg,
		})
	}

	return nil
}

func (r *Registrator) DoScenePansharpen(panScenes []*Scene, mssScenes []*Scene) error {
	for i := 0; i < len(panScenes); i++ {
		fusedTiff := strings.Replace(mssScenes[i].Tiff, "MSS", "FUS", -1)
		err := fusion.Pansharpen(panScenes[i].Tiff, mssScenes[i].Tiff, fusedTiff, fusion.ESRI, 0.1)
		if err != nil {
			return err
		}

		jpg := strings.Replace(fusedTiff, ".tiff", ".jpg", 1)
		GTiffToJPG(fusedTiff, jpg, "FUS", true)
		r.report.Scenes = append(r.report.Scenes, ReportScene{
			TiffData:    fusedTiff,
			Name:        strings.TrimSuffix(filepath.Base(fusedTiff), ".tiff"),
			BrowserData: jpg,
		})
	}
	return nil
}

func (r *Registrator) MergeMSSToBGRNIR(channels []gocv.Mat) (gocv.Mat, error) {
	var rgbirImage gocv.Mat
	if len(channels) != 4 {
		return rgbirImage, fmt.Errorf("mss channels count not match")
	}

	rgbirImage = gocv.NewMat()
	gocv.Merge(channels, &rgbirImage)
	log.Printf("merge mss to bgr nir image, channels: %d, height: %d, width: %d",
		rgbirImage.Channels(), rgbirImage.Rows(), rgbirImage.Cols())

	return rgbirImage, nil

}

func (r *Registrator) RoiScenes() (panScenes []*Scene, mssScenes []*Scene, err error) {
	log.Println("using target scenes")
	for _, target := range r.Params.Targets.Targets {
		y0 := 4 * target.StartLine
		y1 := 4 * target.EndLine
		if y1 > r.PanHeight {
			y1 = r.PanHeight
		}
		scene := &Scene{
			Type:       "PAN",
			Width:      payload.PAN_PIXEL_WIDTH,
			Height:     y1 - y0,
			X:          0,
			Y:          y0,
			SceneIndex: fmt.Sprintf("ROI_%d_%d", target.StartLine, target.EndLine),
		}

		if scene.Height < scene.Width/2 {
			log.Info("scene height too small, skip")
			continue
		}

		name := filepath.Base(r.Params.PanTiffFile)
		name = strings.TrimSuffix(name, ".tiff")
		scene.SceneId = fmt.Sprintf("%s_ROI_%d_%d", name, target.StartLine, target.EndLine)

		mat := r.PanImage.Region(image.Rect(0, y0, payload.PAN_PIXEL_WIDTH, y1))
		if config.GCONFIG.Radiation.PANRemoveHfNoise {
			log.Println("applying hf noise filter on", scene.SceneId)
			matFiltered := rrc.HFNoiseFilter(mat, float64(mat.Cols())*config.GCONFIG.Radiation.HfRadiusRatio)
			mat.Close()
			mat = matFiltered
		}

		if config.GCONFIG.Radiation.SceneMomentMatching {
			rrc.DoMomentMatching(mat)
		}

		scene.Mat = append(scene.Mat, mat)
		panScenes = append(panScenes, scene)
	}

	for _, target := range r.Params.Targets.Targets {
		y0 := target.StartLine
		y1 := target.EndLine
		if y1 > r.MssHeight {
			y1 = r.MssHeight
		}

		scene := &Scene{
			Type:       "MSS",
			Width:      payload.MSS_PIXEL_WIDTH,
			Height:     y1 - y0,
			X:          0,
			Y:          y0,
			SceneIndex: fmt.Sprintf("ROI_%d_%d", target.StartLine, target.EndLine),
		}

		if scene.Height < scene.Width/2 {
			log.Info("scene height too small, skip")
			continue
		}

		name := filepath.Base(r.Params.MssTiffFile)
		name = strings.TrimSuffix(name, ".tiff")
		scene.SceneId = fmt.Sprintf("%s_ROI_%d_%d", name, target.StartLine, target.EndLine)

		for band := 0; band < 4; band++ {
			mat := r.registeredMssImages[band].Region(image.Rect(0, y0, 2336, y1))
			if config.GCONFIG.Radiation.MSSRemoveHfNoise {
				log.Println("applying hf noise filter on", scene.SceneId)
				matFiltered := rrc.HFNoiseFilter(mat, float64(mat.Cols())*config.GCONFIG.Radiation.HfRadiusRatio/4)
				mat.Close()
				mat = matFiltered
			}

			if config.GCONFIG.Radiation.SceneMomentMatching {
				rrc.DoMomentMatching(mat)
			}

			scene.Mat = append(scene.Mat, mat)
		}

		mssScenes = append(mssScenes, scene)
	}

	if len(panScenes) != len(mssScenes) {
		log.Error("pan and mss scenes count not match")
		err = fmt.Errorf("pan and mss scenes count not match")
	}

	return panScenes, mssScenes, err
}
